Skip to content

Commit 1f5c698

Browse files
darth-cydarth-cy
committed
fix
1 parent 38a9ba4 commit 1f5c698

7 files changed

Lines changed: 160 additions & 26 deletions

File tree

‎extensions/native/circuit/src/extension/cuda.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,15 +76,21 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, Native>
7676
inventory.add_executor_chip(fri_reduced_opening);
7777

7878
inventory.next_air::<NativePoseidon2Air<BabyBear, 1>>()?;
79-
let poseidon2 = NativePoseidon2ChipGpu::<1>::new(range_checker.clone(), timestamp_max_bits);
80-
inventory.add_executor_chip(poseidon2);
8179

8280
let hint_air: &HintSpaceProviderAir = inventory.next_air::<HintSpaceProviderAir>()?;
8381
let cpu_chip = Arc::new(HintSpaceProviderChip::new(
8482
hint_air.hint_bus,
8583
range_checker.clone(),
8684
timestamp_max_bits,
8785
));
86+
87+
let poseidon2 = NativePoseidon2ChipGpu::<1>::new_with_hint_space_provider(
88+
range_checker.clone(),
89+
timestamp_max_bits,
90+
cpu_chip.clone(),
91+
);
92+
inventory.add_executor_chip(poseidon2);
93+
8894
let provider_gpu = HintSpaceProviderChipGpu::new(cpu_chip.clone());
8995
inventory.add_periphery_chip(provider_gpu);
9096

‎extensions/native/circuit/src/extension/mod.rs‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -271,15 +271,6 @@ where
271271
);
272272
inventory.add_air(fri_reduced_opening);
273273

274-
let verify_batch = NativePoseidon2Air::<_, 1>::new(
275-
exec_bridge,
276-
memory_bridge,
277-
hint_bridge,
278-
VerifyBatchBus::new(inventory.new_bus_idx()),
279-
Poseidon2Config::default(),
280-
);
281-
inventory.add_air(verify_batch);
282-
283274
let hint_space_provider = HintSpaceProviderAir {
284275
hint_bus: hint_bridge.hint_bus(),
285276
lt_air: IsLtSubAir::new(
@@ -289,6 +280,15 @@ where
289280
};
290281
inventory.add_air(hint_space_provider);
291282

283+
let verify_batch = NativePoseidon2Air::<_, 1>::new(
284+
exec_bridge,
285+
memory_bridge,
286+
hint_bridge,
287+
VerifyBatchBus::new(inventory.new_bus_idx()),
288+
Poseidon2Config::default(),
289+
);
290+
inventory.add_air(verify_batch);
291+
292292
let tower_evaluate = NativeSumcheckAir::new(exec_bridge, memory_bridge, hint_bridge);
293293
inventory.add_air(tower_evaluate);
294294

@@ -365,26 +365,26 @@ where
365365
FriReducedOpeningChip::new(FriReducedOpeningFiller::new(), mem_helper.clone());
366366
inventory.add_executor_chip(fri_reduced_opening);
367367

368-
inventory.next_air::<NativePoseidon2Air<Val<SC>, 1>>()?;
369-
370368
let hint_bus = inventory.airs().system().hint_bridge.hint_bus();
371369
let hint_space_provider = Arc::new(HintSpaceProviderChip::new(
372370
hint_bus,
373371
range_checker.clone(),
374372
timestamp_max_bits,
375373
));
376374

375+
inventory.next_air::<HintSpaceProviderAir>()?;
376+
inventory.add_periphery_chip(hint_space_provider.clone());
377+
378+
inventory.next_air::<NativePoseidon2Air<Val<SC>, 1>>()?;
379+
377380
let poseidon2 = NativePoseidon2Chip::<_, 1>::new(
378381
NativePoseidon2Filler::new(Poseidon2Config::default(), hint_space_provider.clone()),
379382
mem_helper.clone(),
380383
);
381384
inventory.add_executor_chip(poseidon2);
382385

383-
inventory.next_air::<HintSpaceProviderAir>()?;
384-
inventory.add_periphery_chip(hint_space_provider.clone());
385-
386386
let tower_verify = NativeSumcheckChip::new(
387-
NativeSumcheckFiller::new(hint_space_provider),
387+
NativeSumcheckFiller::new(hint_space_provider.clone()),
388388
mem_helper.clone(),
389389
);
390390
inventory.add_executor_chip(tower_verify);

‎extensions/native/circuit/src/poseidon2/chip.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -679,6 +679,10 @@ where
679679
let input_len = ctx[1];
680680
let is_hint = ctx[2].as_canonical_u32() != 0;
681681

682+
683+
// _debug
684+
println!("=> is_hint: {:?}", is_hint);
685+
682686
// Read hint_id from register
683687
let [hint_id]: [F; 1] =
684688
memory_read_native(state.memory.data(), hint_id_register.as_canonical_u32());
@@ -830,6 +834,12 @@ where
830834
);
831835
v
832836
};
837+
838+
// _debug
839+
if is_hint {
840+
println!("multi_observe hint mode: reading nf = {}", n_f);
841+
}
842+
833843
multi_observe_cols.aux_read_enabled[j] = F::ONE;
834844
tracing_write_native_inplace(
835845
state.memory,
@@ -1265,7 +1275,7 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> NativePoseidon2Filler<F, SBOX
12651275

12661276
for j in chunk_start..chunk_end {
12671277
if is_hint {
1268-
// In hint mode: register hint read on the hint bus
1278+
// In hint mode: register hint read on the hint bus.
12691279
let input_idx = multi_observe_cols.curr_len.as_canonical_u32()
12701280
+ (j - chunk_start);
12711281
let val = multi_observe_cols.data[j as usize];

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

Lines changed: 107 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,121 @@
11
use std::{borrow::Borrow, mem::size_of, slice::from_raw_parts, sync::Arc};
22

3-
use derive_new::new;
43
use openvm_circuit::{arch::DenseRecordArena, utils::next_power_of_two_or_zero};
54
use openvm_circuit_primitives::var_range::VariableRangeCheckerChipGPU;
65
use openvm_cuda_backend::{
76
base::DeviceMatrix, chip::get_empty_air_proving_ctx, prover_backend::GpuBackend, types::F,
87
};
98
use openvm_cuda_common::copy::MemCopyH2D;
109
use openvm_stark_backend::{prover::types::AirProvingContext, Chip};
11-
use p3_field::{Field, PrimeField32};
10+
use p3_field::{Field, FieldAlgebra, PrimeField32};
1211

13-
use super::columns::NativePoseidon2Cols;
14-
use crate::cuda_abi::poseidon2_cuda;
12+
use super::columns::{MultiObserveCols, NativePoseidon2Cols};
13+
use crate::{
14+
cuda_abi::poseidon2_cuda,
15+
hint_space_provider::SharedHintSpaceProviderChip,
16+
};
1517

16-
#[derive(new)]
1718
pub struct NativePoseidon2ChipGpu<const SBOX_REGISTERS: usize> {
1819
pub range_checker: Arc<VariableRangeCheckerChipGPU>,
1920
pub timestamp_max_bits: usize,
21+
pub hint_space_provider: Option<SharedHintSpaceProviderChip<F>>,
22+
}
23+
24+
impl<const SBOX_REGISTERS: usize> NativePoseidon2ChipGpu<SBOX_REGISTERS> {
25+
pub fn new(range_checker: Arc<VariableRangeCheckerChipGPU>, timestamp_max_bits: usize) -> Self {
26+
Self {
27+
range_checker,
28+
timestamp_max_bits,
29+
hint_space_provider: None,
30+
}
31+
}
32+
33+
pub fn new_with_hint_space_provider(
34+
range_checker: Arc<VariableRangeCheckerChipGPU>,
35+
timestamp_max_bits: usize,
36+
hint_space_provider: SharedHintSpaceProviderChip<F>,
37+
) -> Self {
38+
Self {
39+
range_checker,
40+
timestamp_max_bits,
41+
hint_space_provider: Some(hint_space_provider),
42+
}
43+
}
44+
45+
/// Scans multi-observe execution records to populate the hint provider with
46+
/// (hint_id, offset, value) triples for hint-mode rows.
47+
fn populate_hint_provider(&self, records: &[u8]) {
48+
let Some(hint_space_provider) = &self.hint_space_provider else {
49+
return;
50+
};
51+
52+
let width = NativePoseidon2Cols::<F, SBOX_REGISTERS>::width();
53+
let record_size = width * size_of::<F>();
54+
if records.len() % record_size != 0 {
55+
return;
56+
}
57+
let height = records.len() / record_size;
58+
59+
let row_slice = unsafe {
60+
let ptr = records.as_ptr() as *const F;
61+
from_raw_parts(ptr, height * width)
62+
};
63+
64+
let mut row_idx = 0;
65+
while row_idx < height {
66+
let start = row_idx * width;
67+
let cols: &NativePoseidon2Cols<F, SBOX_REGISTERS> =
68+
row_slice[start..(start + width)].borrow();
69+
70+
if cols.multi_observe_row.is_one() {
71+
let num_rows = cols.inner.export.as_canonical_u32() as usize;
72+
if num_rows > 1 {
73+
let head_multi_observe_cols: &MultiObserveCols<F> =
74+
cols.specific[..MultiObserveCols::<u8>::width()].borrow();
75+
let is_hint = head_multi_observe_cols.ctx[2] != F::ZERO;
76+
if is_hint {
77+
let hint_id = head_multi_observe_cols.hint_id;
78+
for local_row in 1..num_rows {
79+
let chunk_cols: &NativePoseidon2Cols<F, SBOX_REGISTERS> =
80+
row_slice[(row_idx + local_row) * width
81+
..(row_idx + local_row + 1) * width]
82+
.borrow();
83+
let multi_observe_cols: &MultiObserveCols<F> = chunk_cols.specific
84+
[..MultiObserveCols::<u8>::width()]
85+
.borrow();
86+
87+
let chunk_start = multi_observe_cols.start_idx.as_canonical_u32();
88+
let chunk_end = multi_observe_cols.end_idx.as_canonical_u32();
89+
let curr_len = multi_observe_cols.curr_len.as_canonical_u32();
90+
91+
for j in chunk_start..chunk_end {
92+
let input_idx = curr_len + (j - chunk_start);
93+
let val = multi_observe_cols.data[j as usize];
94+
hint_space_provider.request(
95+
hint_id,
96+
F::from_canonical_u32(input_idx),
97+
val,
98+
);
99+
}
100+
}
101+
}
102+
}
103+
row_idx += num_rows.max(1);
104+
continue;
105+
}
106+
107+
if cols.simple.is_one() {
108+
row_idx += 1;
109+
} else {
110+
let num_non_inside_row = cols.inner.export.as_canonical_u32() as usize;
111+
let non_inside_start = start + (num_non_inside_row - 1) * width;
112+
let last_non_inside_cols: &NativePoseidon2Cols<F, SBOX_REGISTERS> =
113+
row_slice[non_inside_start..(non_inside_start + width)].borrow();
114+
let total_num_row = last_non_inside_cols.inner.export.as_canonical_u32() as usize;
115+
row_idx += total_num_row;
116+
}
117+
}
118+
}
20119
}
21120

22121
impl<const SBOX_REGISTERS: usize> Chip<DenseRecordArena, GpuBackend>
@@ -28,6 +127,9 @@ impl<const SBOX_REGISTERS: usize> Chip<DenseRecordArena, GpuBackend>
28127
return get_empty_air_proving_ctx::<GpuBackend>();
29128
}
30129

130+
// Populate hint space provider from multi-observe records before GPU upload.
131+
self.populate_hint_provider(records);
132+
31133
// For Poseidon2, the records are already the trace rows
32134
// Use the columns width directly
33135
let width = NativePoseidon2Cols::<F, SBOX_REGISTERS>::width();

‎extensions/native/circuit/src/poseidon2/execution.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -640,6 +640,12 @@ unsafe fn execute_multi_observe_e12_impl<
640640
let [v]: [F; 1] = exec_state.vm_read(NATIVE_AS, input_ptr_u32 + input_idx);
641641
v
642642
};
643+
644+
// _debug
645+
if is_hint {
646+
println!("=> n_f: {n_f}");
647+
}
648+
643649
exec_state.vm_write(NATIVE_AS, sponge_ptr_u32 + (j as u32), &[n_f]);
644650
input_idx += 1;
645651
}

‎extensions/native/compiler/src/ir/poseidon.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ impl<C: Config> Builder<C> {
1919
sponge_state: &Array<C, Felt<C::F>>,
2020
input_ptr: Ptr<C::N>,
2121
arr: &Array<C, Felt<C::F>>,
22+
input_len: Usize<C::N>,
2223
hint_id: Option<Var<C::N>>,
2324
) -> Usize<C::N> {
2425
let buffer_size: Var<C::N> = Var::uninit(self);
@@ -32,7 +33,7 @@ impl<C: Config> Builder<C> {
3233
Array::Fixed(_) => {
3334
panic!("Base elements input must be dynamic");
3435
}
35-
Array::Dyn(ptr, len) => {
36+
Array::Dyn(ptr, _) => {
3637
let init_pos: Var<C::N> = Var::uninit(self);
3738
self.assign(&init_pos, input_ptr.address - sponge_ptr.address);
3839

@@ -48,14 +49,23 @@ impl<C: Config> Builder<C> {
4849
// Allocate context array: [init_pos, len, is_hint, reserved]
4950
let ctx = self.dyn_array::<Var<C::N>>(4usize);
5051
self.set(&ctx, 0, init_pos);
51-
self.set(&ctx, 1, len.get_var());
52+
self.set(&ctx, 1, input_len.get_var());
5253
self.set(
5354
&ctx,
5455
2,
5556
if is_hint { C::N::ONE } else { C::N::ZERO },
5657
);
5758
self.set(&ctx, 3, C::N::ZERO);
5859

60+
61+
// _debug
62+
let ctx1 = self.get(&ctx, 1);
63+
let ctx2 = self.get(&ctx, 2);
64+
self.print_debug(777);
65+
self.print_v(ctx1);
66+
self.print_v(ctx2);
67+
68+
5969
self.operations.push(DslIr::Poseidon2MultiObserve(
6070
*sponge_ptr,
6171
ctx.ptr(),

‎extensions/native/recursion/src/challenger/duplex.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ impl<C: Config> DuplexChallengerVariable<C> {
8181
// This is equivalent to calling `observe` multiple times, but more efficient.
8282
pub fn observe_slice_opt(&self, builder: &mut Builder<C>, arr: &Array<C, Felt<C::F>>) {
8383
builder.if_ne(arr.len(), Usize::from(0)).then(|builder| {
84-
let next_pos = builder.poseidon2_multi_observe(&self.sponge_state, self.input_ptr, arr, None);
84+
let next_pos = builder.poseidon2_multi_observe(&self.sponge_state, self.input_ptr, arr, arr.len(), None);
8585

8686
builder.assign(&self.input_ptr, self.io_empty_ptr + next_pos.clone());
8787
builder.if_ne(next_pos, Usize::from(0)).then_or_else(

0 commit comments

Comments
 (0)