finish

Author: kunxian-xia

Cargo.lock

@@ -62,12 +62,43 @@ template <typename T> struct SimplePoseidonSpecificCols {
     MemoryWriteAuxCols<T, CHUNK> write_data_2;
 };
 
+template <typename T> struct MultiObserveCols {
+    T pc;
+    T final_timestamp_increment;
+    T state_ptr;
+    T input_ptr;
+    T init_pos;
+    T len;
+    T input_register_1;
+    T input_register_2;
+    T input_register_3;
+    T output_register;
+    T is_first;
+    T is_last;
+    T curr_len;
+    T start_idx;
+    T end_idx;
+    T aux_after_start[CHUNK];
+    T aux_before_end[CHUNK];
+    T aux_read_enabled[CHUNK];
+    MemoryReadAuxCols<T> read_data[CHUNK];
+    MemoryWriteAuxCols<T, 1> write_data[CHUNK];
+    T data[CHUNK];
+    T should_permute;
+    MemoryWriteAuxCols<T, CHUNK * 2> write_sponge_state;
+    MemoryWriteAuxCols<T, 1> write_final_idx;
+    T final_idx;
+};
+
 template <typename T> constexpr T constexpr_max(T a, T b) { return a > b ? a : b; }
 
 constexpr size_t COL_SPECIFIC_WIDTH = constexpr_max(
     sizeof(TopLevelSpecificCols<uint8_t>),
     constexpr_max(
         sizeof(InsideRowSpecificCols<uint8_t>),
-        sizeof(SimplePoseidonSpecificCols<uint8_t>)
+        constexpr_max(
+            sizeof(SimplePoseidonSpecificCols<uint8_t>),
+            sizeof(MultiObserveCols<uint8_t>)
+        )
     )
 );

extensions/native/circuit/cuda/include/native/poseidon2.cuh

@@ -53,6 +53,9 @@ impl<const SBOX_REGISTERS: usize> Chip<DenseRecordArena, GpuBackend>
                 chunk_start.push(row_idx as u32);
                 if cols.simple.is_one() {
                     row_idx += 1;
+                } else if cols.multi_observe_row.is_one() {
+                    let num_rows = cols.inner.export.as_canonical_u32() as usize;
+                    row_idx += num_rows;
                 } else {
                     let num_non_inside_row = cols.inner.export.as_canonical_u32() as usize;
                     let non_inside_start = start + (num_non_inside_row - 1) * width;

extensions/native/circuit/src/poseidon2/cuda.rs

@@ -8,6 +8,7 @@ repository.workspace = true
 
 [dependencies]
 openvm-stark-backend = { workspace = true }
+openvm-cuda-backend = { workspace = true, optional = true }
 openvm-native-circuit = { workspace = true, features = ["test-utils"] }
 openvm-native-compiler = { workspace = true }
 openvm-native-compiler-derive = { workspace = true }
@@ -58,4 +59,4 @@ parallel = ["openvm-stark-backend/parallel"]
 mimalloc = ["openvm-stark-backend/mimalloc"]
 jemalloc = ["openvm-stark-backend/jemalloc"]
 nightly-features = ["openvm-circuit/nightly-features"]
-cuda = ["openvm-circuit/cuda", "openvm-native-circuit/cuda"]
+cuda = ["openvm-circuit/cuda", "openvm-native-circuit/cuda", "dep:openvm-cuda-backend"]

extensions/native/recursion/Cargo.toml

@@ -6,6 +6,8 @@ use openvm_circuit::{
     },
     utils::{air_test_impl, TestStarkEngine},
 };
+#[cfg(feature = "cuda")]
+use openvm_cuda_backend::engine::GpuBabyBearPoseidon2Engine;
 use openvm_native_circuit::{
     execute_program_with_config, test_native_config, NativeBuilder, NativeConfig,
 };
@@ -211,8 +213,22 @@ fn test_multi_observe() {
     config.system.memory_config.max_access_adapter_n = 16;
 
     let vb = NativeBuilder::default();
+    #[cfg(not(feature = "cuda"))]
     air_test_impl::<BabyBearPoseidon2Engine, _>(fri_params, vb, config, program, vec![], 1, true)
         .unwrap();
+    #[cfg(feature = "cuda")]
+    {
+        air_test_impl::<GpuBabyBearPoseidon2Engine, _>(
+            fri_params,
+            vb,
+            config,
+            program,
+            vec![],
+            1,
+            true,
+        )
+        .unwrap();
+    }
 }
 
 fn build_test_program<C: Config>(builder: &mut Builder<C>) {