From 4c46cbce930524e68c0d16412e87b7cd2e131a2d Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Fri, 9 Aug 2024 12:31:40 -0700 Subject: [PATCH] added interaction scope --- core/src/air/builder.rs | 191 +++++++++++++++++---------------- core/src/lookup/builder.rs | 47 ++++---- core/src/lookup/interaction.rs | 5 + core/src/memory/global.rs | 28 +++-- core/src/memory/program.rs | 17 +-- 5 files changed, 160 insertions(+), 128 deletions(-) diff --git a/core/src/air/builder.rs b/core/src/air/builder.rs index d5d87b2b23..bed8ca72ff 100644 --- a/core/src/air/builder.rs +++ b/core/src/air/builder.rs @@ -17,18 +17,23 @@ use crate::lookup::InteractionKind; use crate::memory::MemoryAccessCols; use crate::{bytes::ByteOpcode, memory::MemoryCols}; +pub enum MessageScope { + Shard, + Global, +} + /// A Builder with the ability to encode the existance of interactions with other AIRs by sending /// and receiving messages. pub trait MessageBuilder { - fn send(&mut self, message: M); + fn send(&mut self, message: M, scope: MessageScope); - fn receive(&mut self, message: M); + fn receive(&mut self, message: M, scope: MessageScope); } impl MessageBuilder for AB { - fn send(&mut self, _message: M) {} + fn send(&mut self, _message: M, _scope: MessageScope) {} - fn receive(&mut self, _message: M) {} + fn receive(&mut self, _message: M, _scope: MessageScope) {} } /// A message builder for which sending and receiving messages is a no-op. @@ -125,19 +130,22 @@ pub trait ByteAirBuilder: BaseAirBuilder { channel: impl Into, multiplicity: impl Into, ) { - self.send(AirInteraction::new( - vec![ - opcode.into(), - a1.into(), - a2.into(), - b.into(), - c.into(), - shard.into(), - channel.into(), - ], - multiplicity.into(), - InteractionKind::Byte, - )); + self.send( + AirInteraction::new( + vec![ + opcode.into(), + a1.into(), + a2.into(), + b.into(), + c.into(), + shard.into(), + channel.into(), + ], + multiplicity.into(), + InteractionKind::Byte, + ), + MessageScope::Global, + ); } /// Receives a byte operation to be processed. @@ -177,19 +185,22 @@ pub trait ByteAirBuilder: BaseAirBuilder { channel: impl Into, multiplicity: impl Into, ) { - self.receive(AirInteraction::new( - vec![ - opcode.into(), - a1.into(), - a2.into(), - b.into(), - c.into(), - shard.into(), - channel.into(), - ], - multiplicity.into(), - InteractionKind::Byte, - )); + self.receive( + AirInteraction::new( + vec![ + opcode.into(), + a1.into(), + a2.into(), + b.into(), + c.into(), + shard.into(), + channel.into(), + ], + multiplicity.into(), + InteractionKind::Byte, + ), + MessageScope::Global, + ); } } @@ -325,11 +336,10 @@ pub trait AluAirBuilder: BaseAirBuilder { .chain(once(nonce.into())) .collect(); - self.send(AirInteraction::new( - values, - multiplicity.into(), - InteractionKind::Alu, - )); + self.send( + AirInteraction::new(values, multiplicity.into(), InteractionKind::Alu), + MessageScope::Global, + ); } /// Receives an ALU operation to be processed. @@ -354,11 +364,10 @@ pub trait AluAirBuilder: BaseAirBuilder { .chain(once(nonce.into())) .collect(); - self.receive(AirInteraction::new( - values, - multiplicity.into(), - InteractionKind::Alu, - )); + self.receive( + AirInteraction::new(values, multiplicity.into(), InteractionKind::Alu), + MessageScope::Global, + ); } /// Sends an syscall operation to be processed (with "ECALL" opcode). @@ -374,19 +383,22 @@ pub trait AluAirBuilder: BaseAirBuilder { arg2: impl Into + Clone, multiplicity: impl Into, ) { - self.send(AirInteraction::new( - vec![ - shard.clone().into(), - channel.clone().into(), - clk.clone().into(), - nonce.clone().into(), - syscall_id.clone().into(), - arg1.clone().into(), - arg2.clone().into(), - ], - multiplicity.into(), - InteractionKind::Syscall, - )); + self.send( + AirInteraction::new( + vec![ + shard.clone().into(), + channel.clone().into(), + clk.clone().into(), + nonce.clone().into(), + syscall_id.clone().into(), + arg1.clone().into(), + arg2.clone().into(), + ], + multiplicity.into(), + InteractionKind::Syscall, + ), + MessageScope::Global, + ); } /// Receives a syscall operation to be processed. @@ -402,19 +414,22 @@ pub trait AluAirBuilder: BaseAirBuilder { arg2: impl Into + Clone, multiplicity: impl Into, ) { - self.receive(AirInteraction::new( - vec![ - shard.clone().into(), - channel.clone().into(), - clk.clone().into(), - nonce.clone().into(), - syscall_id.clone().into(), - arg1.clone().into(), - arg2.clone().into(), - ], - multiplicity.into(), - InteractionKind::Syscall, - )); + self.receive( + AirInteraction::new( + vec![ + shard.clone().into(), + channel.clone().into(), + clk.clone().into(), + nonce.clone().into(), + syscall_id.clone().into(), + arg1.clone().into(), + arg2.clone().into(), + ], + multiplicity.into(), + InteractionKind::Syscall, + ), + MessageScope::Global, + ); } } @@ -466,18 +481,16 @@ pub trait MemoryAirBuilder: BaseAirBuilder { .collect(); // The previous values get sent with multiplicity = 1, for "read". - self.send(AirInteraction::new( - prev_values, - do_check.clone(), - InteractionKind::Memory, - )); + self.send( + AirInteraction::new(prev_values, do_check.clone(), InteractionKind::Memory), + MessageScope::Global, + ); // The current values get "received", i.e. multiplicity = -1 - self.receive(AirInteraction::new( - current_values, - do_check.clone(), - InteractionKind::Memory, - )); + self.receive( + AirInteraction::new(current_values, do_check.clone(), InteractionKind::Memory), + MessageScope::Global, + ); } /// Constraints a memory read or write to a slice of `MemoryAccessCols`. @@ -622,11 +635,10 @@ pub trait ProgramAirBuilder: BaseAirBuilder { .chain(once(shard.into())) .collect(); - self.send(AirInteraction::new( - values, - multiplicity.into(), - InteractionKind::Program, - )); + self.send( + AirInteraction::new(values, multiplicity.into(), InteractionKind::Program), + MessageScope::Global, + ); } /// Receives an instruction. @@ -645,11 +657,10 @@ pub trait ProgramAirBuilder: BaseAirBuilder { .chain(once(shard.into())) .collect(); - self.receive(AirInteraction::new( - values, - multiplicity.into(), - InteractionKind::Program, - )); + self.receive( + AirInteraction::new(values, multiplicity.into(), InteractionKind::Program), + MessageScope::Global, + ); } } @@ -714,12 +725,12 @@ pub trait SP1AirBuilder: } impl<'a, AB: AirBuilder + MessageBuilder, M> MessageBuilder for FilteredAirBuilder<'a, AB> { - fn send(&mut self, message: M) { - self.inner.send(message); + fn send(&mut self, message: M, scope: MessageScope) { + self.inner.send(message, scope); } - fn receive(&mut self, message: M) { - self.inner.receive(message); + fn receive(&mut self, message: M, scope: MessageScope) { + self.inner.receive(message, scope); } } diff --git a/core/src/lookup/builder.rs b/core/src/lookup/builder.rs index a342425d43..04af50c9bc 100644 --- a/core/src/lookup/builder.rs +++ b/core/src/lookup/builder.rs @@ -4,7 +4,7 @@ use p3_matrix::dense::RowMajorMatrix; use p3_uni_stark::{Entry, SymbolicExpression, SymbolicVariable}; use crate::{ - air::{AirInteraction, MessageBuilder}, + air::{AirInteraction, MessageBuilder, MessageScope}, stark::PROOF_MAX_NUM_PVS, }; @@ -91,7 +91,7 @@ impl PairBuilder for InteractionBuilder { } impl MessageBuilder>> for InteractionBuilder { - fn send(&mut self, message: AirInteraction>) { + fn send(&mut self, message: AirInteraction>, scope: MessageScope) { let values = message .values .into_iter() @@ -101,10 +101,10 @@ impl MessageBuilder>> for Interac let multiplicity = symbolic_to_virtual_pair(&message.multiplicity); self.sends - .push(Interaction::new(values, multiplicity, message.kind)); + .push(Interaction::new(values, multiplicity, message.kind, scope)); } - fn receive(&mut self, message: AirInteraction>) { + fn receive(&mut self, message: AirInteraction>, scope: MessageScope) { let values = message .values .into_iter() @@ -114,7 +114,7 @@ impl MessageBuilder>> for Interac let multiplicity = symbolic_to_virtual_pair(&message.multiplicity); self.receives - .push(Interaction::new(values, multiplicity, message.kind)); + .push(Interaction::new(values, multiplicity, message.kind, scope)); } } @@ -250,22 +250,27 @@ mod tests { let y = local[1]; let z = local[2]; - builder.send(AirInteraction::new( - vec![x.into(), y.into()], - AB::F::from_canonical_u32(3).into(), - InteractionKind::Alu, - )); - builder.send(AirInteraction::new( - vec![x + y, z.into()], - AB::F::from_canonical_u32(5).into(), - InteractionKind::Alu, - )); - - builder.receive(AirInteraction::new( - vec![x.into()], - y.into(), - InteractionKind::Byte, - )); + builder.send( + AirInteraction::new( + vec![x.into(), y.into()], + AB::F::from_canonical_u32(3).into(), + InteractionKind::Alu, + ), + MessageScope::Global, + ); + builder.send( + AirInteraction::new( + vec![x + y, z.into()], + AB::F::from_canonical_u32(5).into(), + InteractionKind::Alu, + ), + MessageScope::Global, + ); + + builder.receive( + AirInteraction::new(vec![x.into()], y.into(), InteractionKind::Byte), + MessageScope::Global, + ); } } diff --git a/core/src/lookup/interaction.rs b/core/src/lookup/interaction.rs index 1c20938cc4..7a1323c4da 100644 --- a/core/src/lookup/interaction.rs +++ b/core/src/lookup/interaction.rs @@ -4,11 +4,14 @@ use core::fmt::Display; use p3_air::VirtualPairCol; use p3_field::Field; +use crate::air::MessageScope; + /// An interaction for a lookup or a permutation argument. pub struct Interaction { pub values: Vec>, pub multiplicity: VirtualPairCol, pub kind: InteractionKind, + pub scope: MessageScope, } /// The type of interaction for a lookup argument. @@ -60,11 +63,13 @@ impl Interaction { values: Vec>, multiplicity: VirtualPairCol, kind: InteractionKind, + scope: MessageScope, ) -> Self { Self { values, multiplicity, kind, + scope, } } diff --git a/core/src/memory/global.rs b/core/src/memory/global.rs index 8bb93618ba..558b35a198 100644 --- a/core/src/memory/global.rs +++ b/core/src/memory/global.rs @@ -10,7 +10,7 @@ use p3_matrix::Matrix; use sp1_derive::AlignedBorrow; use super::{MemoryChipType, MemoryInitializeFinalizeEvent}; -use crate::air::{AirInteraction, BaseAirBuilder, PublicValues, SP1AirBuilder, Word}; +use crate::air::{AirInteraction, BaseAirBuilder, MessageScope, PublicValues, SP1AirBuilder, Word}; use crate::air::{MachineAir, SP1_PROOF_NUM_PV_ELTS}; use crate::operations::{AssertLtColsBits, BabyBearBitDecomposition, IsZeroOperation}; use crate::runtime::{ExecutionRecord, Program}; @@ -219,11 +219,14 @@ where if self.kind == MemoryChipType::Initialize { let mut values = vec![AB::Expr::zero(), AB::Expr::zero(), local.addr.into()]; values.extend(value.map(Into::into)); - builder.receive(AirInteraction::new( - values, - local.is_real.into(), - crate::lookup::InteractionKind::Memory, - )); + builder.receive( + AirInteraction::new( + values, + local.is_real.into(), + crate::lookup::InteractionKind::Memory, + ), + MessageScope::Global, + ); } else { let mut values = vec![ local.shard.into(), @@ -231,11 +234,14 @@ where local.addr.into(), ]; values.extend(value); - builder.send(AirInteraction::new( - values, - local.is_real.into(), - crate::lookup::InteractionKind::Memory, - )); + builder.send( + AirInteraction::new( + values, + local.is_real.into(), + crate::lookup::InteractionKind::Memory, + ), + MessageScope::Global, + ); } // Canonically decompose the address into bits so we can do comparisons. diff --git a/core/src/memory/program.rs b/core/src/memory/program.rs index 3adf5aec1f..4e549d4424 100644 --- a/core/src/memory/program.rs +++ b/core/src/memory/program.rs @@ -8,7 +8,9 @@ use p3_matrix::Matrix; use sp1_derive::AlignedBorrow; -use crate::air::{AirInteraction, PublicValues, SP1AirBuilder, SP1_PROOF_NUM_PV_ELTS}; +use crate::air::{ + AirInteraction, MessageScope, PublicValues, SP1AirBuilder, SP1_PROOF_NUM_PV_ELTS, +}; use crate::air::{MachineAir, Word}; use crate::operations::IsZeroOperation; use crate::runtime::{ExecutionRecord, Program}; @@ -193,10 +195,13 @@ where let mut values = vec![AB::Expr::zero(), AB::Expr::zero(), prep_local.addr.into()]; values.extend(prep_local.value.map(Into::into)); - builder.receive(AirInteraction::new( - values, - mult_local.multiplicity.into(), - crate::lookup::InteractionKind::Memory, - )); + builder.receive( + AirInteraction::new( + values, + mult_local.multiplicity.into(), + crate::lookup::InteractionKind::Memory, + ), + MessageScope::Global, + ); } }