Skip to content

Commit

Permalink
feat: byte multiplicity channel (#800)
Browse files Browse the repository at this point in the history
  • Loading branch information
tamirhemo authored and erabinov committed May 24, 2024
1 parent b4a159b commit e0a941d
Show file tree
Hide file tree
Showing 69 changed files with 1,521 additions and 480 deletions.
64 changes: 61 additions & 3 deletions core/src/air/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,16 +84,27 @@ pub trait BaseAirBuilder: AirBuilder + MessageBuilder<AirInteraction<Self::Expr>
/// A trait which contains methods for byte interactions in an AIR.
pub trait ByteAirBuilder: BaseAirBuilder {
/// Sends a byte operation to be processed.
#[allow(clippy::too_many_arguments)]
fn send_byte(
&mut self,
opcode: impl Into<Self::Expr>,
a: impl Into<Self::Expr>,
b: impl Into<Self::Expr>,
c: impl Into<Self::Expr>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.send_byte_pair(opcode, a, Self::Expr::zero(), b, c, shard, multiplicity)
self.send_byte_pair(
opcode,
a,
Self::Expr::zero(),
b,
c,
shard,
channel,
multiplicity,
)
}

/// Sends a byte operation with two outputs to be processed.
Expand All @@ -106,6 +117,7 @@ pub trait ByteAirBuilder: BaseAirBuilder {
b: impl Into<Self::Expr>,
c: impl Into<Self::Expr>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.send(AirInteraction::new(
Expand All @@ -116,23 +128,35 @@ pub trait ByteAirBuilder: BaseAirBuilder {
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
));
}

/// Receives a byte operation to be processed.
#[allow(clippy::too_many_arguments)]
fn receive_byte(
&mut self,
opcode: impl Into<Self::Expr>,
a: impl Into<Self::Expr>,
b: impl Into<Self::Expr>,
c: impl Into<Self::Expr>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.receive_byte_pair(opcode, a, Self::Expr::zero(), b, c, shard, multiplicity)
self.receive_byte_pair(
opcode,
a,
Self::Expr::zero(),
b,
c,
shard,
channel,
multiplicity,
)
}

/// Receives a byte operation with two outputs to be processed.
Expand All @@ -145,6 +169,7 @@ pub trait ByteAirBuilder: BaseAirBuilder {
b: impl Into<Self::Expr>,
c: impl Into<Self::Expr>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
self.receive(AirInteraction::new(
Expand All @@ -155,6 +180,7 @@ pub trait ByteAirBuilder: BaseAirBuilder {
b.into(),
c.into(),
shard.into(),
channel.into(),
],
multiplicity.into(),
InteractionKind::Byte,
Expand Down Expand Up @@ -219,6 +245,7 @@ pub trait WordAirBuilder: ByteAirBuilder {
&mut self,
input: &[impl Into<Self::Expr> + Clone],
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
mult: impl Into<Self::Expr> + Clone,
) {
let mut index = 0;
Expand All @@ -229,6 +256,7 @@ pub trait WordAirBuilder: ByteAirBuilder {
input[index].clone(),
input[index + 1].clone(),
shard.clone(),
channel.clone(),
mult.clone(),
);
index += 2;
Expand All @@ -240,6 +268,7 @@ pub trait WordAirBuilder: ByteAirBuilder {
input[index].clone(),
Self::Expr::zero(),
shard.clone(),
channel.clone(),
mult.clone(),
);
}
Expand All @@ -250,6 +279,7 @@ pub trait WordAirBuilder: ByteAirBuilder {
&mut self,
input: &[impl Into<Self::Expr> + Copy],
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
mult: impl Into<Self::Expr> + Clone,
) {
input.iter().for_each(|limb| {
Expand All @@ -259,6 +289,7 @@ pub trait WordAirBuilder: ByteAirBuilder {
Self::Expr::zero(),
Self::Expr::zero(),
shard.clone(),
channel.clone(),
mult.clone(),
);
});
Expand All @@ -268,20 +299,23 @@ pub trait WordAirBuilder: ByteAirBuilder {
/// A trait which contains methods related to ALU interactions in an AIR.
pub trait AluAirBuilder: BaseAirBuilder {
/// Sends an ALU operation to be processed.
#[allow(clippy::too_many_arguments)]
fn send_alu(
&mut self,
opcode: impl Into<Self::Expr>,
a: Word<impl Into<Self::Expr>>,
b: Word<impl Into<Self::Expr>>,
c: Word<impl Into<Self::Expr>>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
let values = once(opcode.into())
.chain(a.0.into_iter().map(Into::into))
.chain(b.0.into_iter().map(Into::into))
.chain(c.0.into_iter().map(Into::into))
.chain(once(shard.into()))
.chain(once(channel.into()))
.collect();

self.send(AirInteraction::new(
Expand All @@ -292,20 +326,23 @@ pub trait AluAirBuilder: BaseAirBuilder {
}

/// Receives an ALU operation to be processed.
#[allow(clippy::too_many_arguments)]
fn receive_alu(
&mut self,
opcode: impl Into<Self::Expr>,
a: Word<impl Into<Self::Expr>>,
b: Word<impl Into<Self::Expr>>,
c: Word<impl Into<Self::Expr>>,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
multiplicity: impl Into<Self::Expr>,
) {
let values = once(opcode.into())
.chain(a.0.into_iter().map(Into::into))
.chain(b.0.into_iter().map(Into::into))
.chain(c.0.into_iter().map(Into::into))
.chain(once(shard.into()))
.chain(once(channel.into()))
.collect();

self.receive(AirInteraction::new(
Expand All @@ -316,9 +353,11 @@ pub trait AluAirBuilder: BaseAirBuilder {
}

/// Sends an syscall operation to be processed (with "ECALL" opcode).
#[allow(clippy::too_many_arguments)]
fn send_syscall(
&mut self,
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
clk: impl Into<Self::Expr> + Clone,
syscall_id: impl Into<Self::Expr> + Clone,
arg1: impl Into<Self::Expr> + Clone,
Expand All @@ -328,6 +367,7 @@ pub trait AluAirBuilder: BaseAirBuilder {
self.send(AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
Expand All @@ -339,9 +379,11 @@ pub trait AluAirBuilder: BaseAirBuilder {
}

/// Receives a syscall operation to be processed.
#[allow(clippy::too_many_arguments)]
fn receive_syscall(
&mut self,
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
clk: impl Into<Self::Expr> + Clone,
syscall_id: impl Into<Self::Expr> + Clone,
arg1: impl Into<Self::Expr> + Clone,
Expand All @@ -351,6 +393,7 @@ pub trait AluAirBuilder: BaseAirBuilder {
self.receive(AirInteraction::new(
vec![
shard.clone().into(),
channel.clone().into(),
clk.clone().into(),
syscall_id.clone().into(),
arg1.clone().into(),
Expand All @@ -371,20 +414,28 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
fn eval_memory_access<E: Into<Self::Expr> + Clone>(
&mut self,
shard: impl Into<Self::Expr>,
channel: impl Into<Self::Expr>,
clk: impl Into<Self::Expr>,
addr: impl Into<Self::Expr>,
memory_access: &impl MemoryCols<E>,
do_check: impl Into<Self::Expr>,
) {
let do_check: Self::Expr = do_check.into();
let shard: Self::Expr = shard.into();
let channel: Self::Expr = channel.into();
let clk: Self::Expr = clk.into();
let mem_access = memory_access.access();

self.assert_bool(do_check.clone());

// Verify that the current memory access time is greater than the previous's.
self.eval_memory_access_timestamp(mem_access, do_check.clone(), shard.clone(), clk.clone());
self.eval_memory_access_timestamp(
mem_access,
do_check.clone(),
shard.clone(),
channel,
clk.clone(),
);

// Add to the memory argument.
let addr = addr.into();
Expand Down Expand Up @@ -420,6 +471,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
fn eval_memory_access_slice<E: Into<Self::Expr> + Copy>(
&mut self,
shard: impl Into<Self::Expr> + Copy,
channel: impl Into<Self::Expr> + Clone,
clk: impl Into<Self::Expr> + Clone,
initial_addr: impl Into<Self::Expr> + Clone,
memory_access_slice: &[impl MemoryCols<E>],
Expand All @@ -428,6 +480,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
for (i, access_slice) in memory_access_slice.iter().enumerate() {
self.eval_memory_access(
shard,
channel.clone(),
clk.clone(),
initial_addr.clone().into() + Self::Expr::from_canonical_usize(i * 4),
access_slice,
Expand All @@ -447,6 +500,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
mem_access: &MemoryAccessCols<impl Into<Self::Expr> + Clone>,
do_check: impl Into<Self::Expr>,
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
clk: impl Into<Self::Expr>,
) {
let do_check: Self::Expr = do_check.into();
Expand Down Expand Up @@ -487,6 +541,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
mem_access.diff_16bit_limb.clone(),
mem_access.diff_8bit_limb.clone(),
shard.clone(),
channel.clone(),
do_check,
);
}
Expand All @@ -503,6 +558,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
limb_16: impl Into<Self::Expr> + Clone,
limb_8: impl Into<Self::Expr> + Clone,
shard: impl Into<Self::Expr> + Clone,
channel: impl Into<Self::Expr> + Clone,
do_check: impl Into<Self::Expr> + Clone,
) {
// Verify that value = limb_16 + limb_8 * 2^16.
Expand All @@ -519,6 +575,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
Self::Expr::zero(),
Self::Expr::zero(),
shard.clone(),
channel.clone(),
do_check.clone(),
);

Expand All @@ -528,6 +585,7 @@ pub trait MemoryAirBuilder: BaseAirBuilder {
Self::Expr::zero(),
limb_8,
shard.clone(),
channel.clone(),
do_check,
)
}
Expand Down
24 changes: 19 additions & 5 deletions core/src/alu/add_sub/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@ pub struct AddSubCols<T> {
/// The shard number, used for byte lookup table.
pub shard: T,

/// The channel number, used for byte lookup table.
pub channel: T,

/// Instance of `AddOperation` to handle addition logic in `AddSubChip`'s ALU operations.
/// It's result will be `a` for the add operation and `b` for the sub operation.
pub add_operation: AddOperation<T>,
Expand Down Expand Up @@ -88,14 +91,20 @@ impl<F: PrimeField> MachineAir<F> for AddSubChip {
let cols: &mut AddSubCols<F> = row.as_mut_slice().borrow_mut();
let is_add = event.opcode == Opcode::ADD;
cols.shard = F::from_canonical_u32(event.shard);
cols.channel = F::from_canonical_u32(event.channel);
cols.is_add = F::from_bool(is_add);
cols.is_sub = F::from_bool(!is_add);

let operand_1 = if is_add { event.b } else { event.a };
let operand_2 = event.c;

cols.add_operation
.populate(&mut record, event.shard, operand_1, operand_2);
cols.add_operation.populate(
&mut record,
event.shard,
event.channel,
operand_1,
operand_2,
);
cols.operand_1 = Word::from(operand_1);
cols.operand_2 = Word::from(operand_2);
row
Expand Down Expand Up @@ -150,6 +159,7 @@ where
local.operand_2,
local.add_operation,
local.shard,
local.channel,
local.is_add + local.is_sub,
);

Expand All @@ -161,6 +171,7 @@ where
local.operand_1,
local.operand_2,
local.shard,
local.channel,
local.is_add,
);

Expand All @@ -171,6 +182,7 @@ where
local.add_operation.value,
local.operand_2,
local.shard,
local.channel,
local.is_sub,
);

Expand Down Expand Up @@ -203,7 +215,7 @@ mod tests {
#[test]
fn generate_trace() {
let mut shard = ExecutionRecord::default();
shard.add_events = vec![AluEvent::new(0, 0, Opcode::ADD, 14, 8, 6)];
shard.add_events = vec![AluEvent::new(0, 0, 0, Opcode::ADD, 14, 8, 6)];
let chip = AddSubChip::default();
let trace: RowMajorMatrix<BabyBear> =
chip.generate_trace(&shard, &mut ExecutionRecord::default());
Expand All @@ -216,25 +228,27 @@ mod tests {
let mut challenger = config.challenger();

let mut shard = ExecutionRecord::default();
for _ in 0..1000 {
for i in 0..1000 {
let operand_1 = thread_rng().gen_range(0..u32::MAX);
let operand_2 = thread_rng().gen_range(0..u32::MAX);
let result = operand_1.wrapping_add(operand_2);
shard.add_events.push(AluEvent::new(
0,
i % 2,
0,
Opcode::ADD,
result,
operand_1,
operand_2,
));
}
for _ in 0..1000 {
for i in 0..1000 {
let operand_1 = thread_rng().gen_range(0..u32::MAX);
let operand_2 = thread_rng().gen_range(0..u32::MAX);
let result = operand_1.wrapping_sub(operand_2);
shard.add_events.push(AluEvent::new(
0,
i % 2,
0,
Opcode::SUB,
result,
Expand Down
Loading

0 comments on commit e0a941d

Please sign in to comment.