removed last eq lookup in comparison instructions (#471)

a16z · Sep 30, 2024 · 6dac58c · 6dac58c
1 parent 7625b1d
commit 6dac58c
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/jolt-core/src/jolt/instruction/bge.rs b/jolt-core/src/jolt/instruction/bge.rs
@@ -38,7 +38,7 @@ impl<const WORD_SIZE: usize> JoltInstruction for BGEInstruction<WORD_SIZE> {
             (Box::new(LeftMSBSubtable::new()), SubtableIndices::from(0)),
             (Box::new(RightMSBSubtable::new()), SubtableIndices::from(0)),
             (Box::new(LtuSubtable::new()), SubtableIndices::from(1..C)),
-            (Box::new(EqSubtable::new()), SubtableIndices::from(1..C)),
+            (Box::new(EqSubtable::new()), SubtableIndices::from(1..C - 1)),
             (Box::new(LtAbsSubtable::new()), SubtableIndices::from(0)),
             (Box::new(EqAbsSubtable::new()), SubtableIndices::from(0)),
         ]

diff --git a/jolt-core/src/jolt/instruction/bgeu.rs b/jolt-core/src/jolt/instruction/bgeu.rs
@@ -33,7 +33,7 @@ impl<const WORD_SIZE: usize> JoltInstruction for BGEUInstruction<WORD_SIZE> {
     ) -> Vec<(Box<dyn LassoSubtable<F>>, SubtableIndices)> {
         vec![
             (Box::new(LtuSubtable::new()), SubtableIndices::from(0..C)),
-            (Box::new(EqSubtable::new()), SubtableIndices::from(0..C)),
+            (Box::new(EqSubtable::new()), SubtableIndices::from(0..C - 1)),
         ]
     }
 

diff --git a/jolt-core/src/jolt/instruction/slt.rs b/jolt-core/src/jolt/instruction/slt.rs
@@ -35,10 +35,12 @@ impl<const WORD_SIZE: usize> JoltInstruction for SLTInstruction<WORD_SIZE> {
         // Accumulator for EQ(x_{<s}, y_{<s})
         let mut eq_prod = eq_abs[0];
 
-        for (ltu_i, eq_i) in ltu.iter().zip(eq) {
-            ltu_sum += *ltu_i * eq_prod;
-            eq_prod *= *eq_i;
+        for i in 0..C - 2 {
+            ltu_sum += ltu[i] * eq_prod;
+            eq_prod *= eq[i];
         }
+        // Do not need to update `eq_prod` for the last iteration
+        ltu_sum += ltu[C - 2] * eq_prod;
 
         // x_s * (1 - y_s) + EQ(x_s, y_s) * LTU(x_{<s}, y_{<s})
         left_msb[0] * (F::one() - right_msb[0])
@@ -59,7 +61,7 @@ impl<const WORD_SIZE: usize> JoltInstruction for SLTInstruction<WORD_SIZE> {
             (Box::new(LeftMSBSubtable::new()), SubtableIndices::from(0)),
             (Box::new(RightMSBSubtable::new()), SubtableIndices::from(0)),
             (Box::new(LtuSubtable::new()), SubtableIndices::from(1..C)),
-            (Box::new(EqSubtable::new()), SubtableIndices::from(1..C)),
+            (Box::new(EqSubtable::new()), SubtableIndices::from(1..C - 1)),
             (Box::new(LtAbsSubtable::new()), SubtableIndices::from(0)),
             (Box::new(EqAbsSubtable::new()), SubtableIndices::from(0)),
         ]

diff --git a/jolt-core/src/jolt/instruction/sltu.rs b/jolt-core/src/jolt/instruction/sltu.rs
@@ -28,11 +28,12 @@ impl<const WORD_SIZE: usize> JoltInstruction for SLTUInstruction<WORD_SIZE> {
         let mut sum = F::zero();
         let mut eq_prod = F::one();
 
-        for i in 0..C {
+        for i in 0..C - 1 {
             sum += ltu[i] * eq_prod;
             eq_prod *= eq[i];
         }
-        sum
+        // Do not need to update `eq_prod` for the last iteration
+        sum + ltu[C - 1] * eq_prod
     }
 
     fn g_poly_degree(&self, C: usize) -> usize {
@@ -46,7 +47,7 @@ impl<const WORD_SIZE: usize> JoltInstruction for SLTUInstruction<WORD_SIZE> {
     ) -> Vec<(Box<dyn LassoSubtable<F>>, SubtableIndices)> {
         vec![
             (Box::new(LtuSubtable::new()), SubtableIndices::from(0..C)),
-            (Box::new(EqSubtable::new()), SubtableIndices::from(0..C)),
+            (Box::new(EqSubtable::new()), SubtableIndices::from(0..C - 1)),
         ]
     }