chore(cuda): update poseidon2 kernel to use memory manager (#2108)

Co-authored-by: stephenh-axiom-xyz <[email protected]>

Author: gaxiom

crates/vm/cuda/src/system/poseidon2.cu

@@ -74,34 +74,30 @@ extern "C" int _system_poseidon2_tracegen(
     return cudaGetLastError();
 }
 
-// Reduces the records, removing duplicates and storing the number of times
-// each occurs in d_counts. The number of records after reduction is stored
-// into host pointer num_records.
-extern "C" int _system_poseidon2_deduplicate_records(
+// Prepares d_num_records for use with sort reduce and stores the temporary buffer
+// size necessary for both cub functions (i.e. sort and reduce).
+extern "C" int _system_poseidon2_deduplicate_records_get_temp_bytes(
     Fp *d_records,
     uint32_t *d_counts,
-    size_t *num_records
+    size_t num_records,
+    size_t *d_num_records,
+    size_t *h_temp_bytes_out
 ) {
-    auto [grid, block] = kernel_launch_params(*num_records);
+    auto [grid, block] = kernel_launch_params(num_records);
     FpArray<16> *d_records_fp16 = reinterpret_cast<FpArray<16> *>(d_records);
-    size_t *d_num_records;
 
     // We want to sort and reduce the raw records, keeping track of how many
-    // each occurs in d_counts. To prepare for reduce we need to a) allocate
-    // d_num_records, b) fill d_counts with 1s, and c) group keys together
-    // using sort.
-    cudaMallocAsync(&d_num_records, sizeof(size_t), cudaStreamPerThread);
-    cudaMemcpyAsync(
-        d_num_records, num_records, sizeof(size_t), cudaMemcpyHostToDevice, cudaStreamPerThread
-    );
-    fill_buffer<uint32_t><<<grid, block, 0, cudaStreamPerThread>>>(d_counts, 1, *num_records);
+    // each occurs in d_counts. To prepare for reduce we need to a) fill
+    // d_counts with 1s, and b) group keys together using sort. Note we do
+    // b) in the kernel below.
+    fill_buffer<uint32_t><<<grid, block>>>(d_counts, 1, num_records);
 
     size_t sort_storage_bytes = 0;
     cub::DeviceMergeSort::SortKeys(
         nullptr,
         sort_storage_bytes,
         d_records_fp16,
-        *num_records,
+        num_records,
         Fp16CompareOp(),
         cudaStreamPerThread
     );
@@ -116,13 +112,27 @@ extern "C" int _system_poseidon2_deduplicate_records(
         d_counts,
         d_num_records,
         std::plus(),
-        *num_records,
+        num_records,
         cudaStreamPerThread
     );
 
-    size_t temp_storage_bytes = std::max(sort_storage_bytes, reduce_storage_bytes);
-    void *d_temp_storage = nullptr;
-    cudaMallocAsync(&d_temp_storage, temp_storage_bytes, cudaStreamPerThread);
+    *h_temp_bytes_out = std::max(sort_storage_bytes, reduce_storage_bytes);
+    return cudaGetLastError();
+}
+
+// Reduces the records, removing duplicates and storing the number of times
+// each occurs in d_counts. The number of records after reduction is stored
+// into host pointer num_records. The value of temp_storage_bytes should be
+// computed using _system_poseidon2_deduplicate_records_get_temp_bytes.
+extern "C" int _system_poseidon2_deduplicate_records(
+    Fp *d_records,
+    uint32_t *d_counts,
+    size_t num_records,
+    size_t *d_num_records,
+    void *d_temp_storage,
+    size_t temp_storage_bytes
+) {
+    FpArray<16> *d_records_fp16 = reinterpret_cast<FpArray<16> *>(d_records);
 
     // TODO: We currently can't use DeviceRadixSort since each key is 64 bytes
     // which causes Fp16Decomposer usage to exceed shared memory. We need to
@@ -131,7 +141,7 @@ extern "C" int _system_poseidon2_deduplicate_records(
         d_temp_storage,
         temp_storage_bytes,
         d_records_fp16,
-        *num_records,
+        num_records,
         Fp16CompareOp(),
         cudaStreamPerThread
     );
@@ -148,14 +158,9 @@ extern "C" int _system_poseidon2_deduplicate_records(
         d_counts,
         d_num_records,
         std::plus(),
-        *num_records,
+        num_records,
         cudaStreamPerThread
     );
 
-    cudaMemcpyAsync(
-        num_records, d_num_records, sizeof(size_t), cudaMemcpyDeviceToHost, cudaStreamPerThread
-    );
-    cudaFreeAsync(d_num_records, cudaStreamPerThread);
-    cudaFreeAsync(d_temp_storage, cudaStreamPerThread);
     return cudaGetLastError();
 }

crates/vm/src/cuda_abi.rs

@@ -127,10 +127,21 @@ pub mod poseidon2 {
             sbox_regs: usize,
         ) -> i32;
 
+        fn _system_poseidon2_deduplicate_records_get_temp_bytes(
+            d_records: *mut F,
+            d_counts: *mut u32,
+            num_records: usize,
+            d_num_records: *mut usize,
+            h_temp_bytes_out: *mut usize,
+        ) -> i32;
+
         fn _system_poseidon2_deduplicate_records(
             d_records: *mut F,
             d_counts: *mut u32,
-            num_records: *mut usize,
+            num_records: usize,
+            d_num_records: *mut usize,
+            d_temp_storage: *mut std::ffi::c_void,
+            temp_storage_bytes: usize,
         ) -> i32;
     }
 
@@ -154,15 +165,37 @@ pub mod poseidon2 {
         ))
     }
 
+    pub unsafe fn deduplicate_records_get_temp_bytes(
+        d_records: &DeviceBuffer<F>,
+        d_counts: &DeviceBuffer<u32>,
+        num_records: usize,
+        d_num_records: &DeviceBuffer<usize>,
+        h_temp_bytes_out: &mut usize,
+    ) -> Result<(), CudaError> {
+        CudaError::from_result(_system_poseidon2_deduplicate_records_get_temp_bytes(
+            d_records.as_mut_ptr(),
+            d_counts.as_mut_ptr(),
+            num_records,
+            d_num_records.as_mut_ptr(),
+            h_temp_bytes_out,
+        ))
+    }
+
     pub unsafe fn deduplicate_records(
         d_records: &DeviceBuffer<F>,
         d_counts: &DeviceBuffer<u32>,
-        num_records: &mut usize,
+        num_records: usize,
+        d_num_records: &DeviceBuffer<usize>,
+        d_temp_storage: &DeviceBuffer<u8>,
+        temp_storage_bytes: usize,
     ) -> Result<(), CudaError> {
         CudaError::from_result(_system_poseidon2_deduplicate_records(
             d_records.as_mut_ptr(),
             d_counts.as_mut_ptr(),
-            num_records as *mut usize,
+            num_records,
+            d_num_records.as_mut_ptr(),
+            d_temp_storage.as_mut_raw_ptr(),
+            temp_storage_bytes,
         ))
     }
 }

crates/vm/src/system/cuda/poseidon2.rs

@@ -6,7 +6,10 @@ use openvm_circuit::{
     system::poseidon2::columns::Poseidon2PeripheryCols, utils::next_power_of_two_or_zero,
 };
 use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, prover_backend::GpuBackend};
-use openvm_cuda_common::{copy::MemCopyD2H, d_buffer::DeviceBuffer};
+use openvm_cuda_common::{
+    copy::{MemCopyD2H, MemCopyH2D},
+    d_buffer::DeviceBuffer,
+};
 use openvm_stark_backend::{
     prover::{hal::MatrixDimensions, types::AirProvingContext},
     Chip,
@@ -60,8 +63,27 @@ impl<RA, const SBOX_REGISTERS: usize> Chip<RA, GpuBackend> for Poseidon2ChipGPU<
         let mut num_records = self.idx.to_host().unwrap()[0] as usize;
         let counts = DeviceBuffer::<u32>::with_capacity(num_records);
         unsafe {
-            poseidon2::deduplicate_records(&self.records, &counts, &mut num_records)
-                .expect("Failed to deduplicate records");
+            let d_num_records = [num_records].to_device().unwrap();
+            let mut temp_bytes = 0;
+            poseidon2::deduplicate_records_get_temp_bytes(
+                &self.records,
+                &counts,
+                num_records,
+                &d_num_records,
+                &mut temp_bytes,
+            )
+            .expect("Failed to get temp bytes");
+            let d_temp_storage = DeviceBuffer::<u8>::with_capacity(temp_bytes);
+            poseidon2::deduplicate_records(
+                &self.records,
+                &counts,
+                num_records,
+                &d_num_records,
+                &d_temp_storage,
+                temp_bytes,
+            )
+            .expect("Failed to deduplicate records");
+            num_records = *d_num_records.to_host().unwrap().first().unwrap();
         }
         #[cfg(feature = "metrics")]
         self.current_trace_height