From f30431b09faaec05b77e25c44507b4635df7ade4 Mon Sep 17 00:00:00 2001
From: Alon Haramati <alonh@starkware.co>
Date: Tue, 25 Jun 2024 11:02:07 +0300
Subject: [PATCH] Batch inverse lookup column denominators.

---
 crates/prover/src/core/queries.rs             |  2 --
 crates/prover/src/examples/fibonacci/mod.rs   |  1 -
 .../src/examples/wide_fibonacci/simd.rs       |  1 -
 .../src/examples/wide_fibonacci/trace_gen.rs  | 26 ++++++++++++-------
 4 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/crates/prover/src/core/queries.rs b/crates/prover/src/core/queries.rs
index f2860fe5c..09954d34a 100644
--- a/crates/prover/src/core/queries.rs
+++ b/crates/prover/src/core/queries.rs
@@ -8,8 +8,6 @@ use super::circle::Coset;
 use super::poly::circle::CircleDomain;
 use super::utils::bit_reverse_index;
 
-// TODO(AlonH): Move file to fri directory.
-
 pub const UPPER_BOUND_QUERY_BYTES: usize = 4;
 
 /// An ordered set of query indices over a bit reversed [CircleDomain].
diff --git a/crates/prover/src/examples/fibonacci/mod.rs b/crates/prover/src/examples/fibonacci/mod.rs
index cf2862162..b00bf006a 100644
--- a/crates/prover/src/examples/fibonacci/mod.rs
+++ b/crates/prover/src/examples/fibonacci/mod.rs
@@ -32,7 +32,6 @@ impl Fibonacci {
     pub fn get_trace(&self) -> CpuCircleEvaluation<BaseField, BitReversedOrder> {
         // Trace.
         let trace_domain = CanonicCoset::new(self.air.component.log_size);
-        // TODO(AlonH): Consider using Vec::new instead of Vec::with_capacity throughout file.
         let mut trace = Vec::with_capacity(trace_domain.size());
 
         // Fill trace with fibonacci squared.
diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs
index d2759eb4c..348c77740 100644
--- a/crates/prover/src/examples/wide_fibonacci/simd.rs
+++ b/crates/prover/src/examples/wide_fibonacci/simd.rs
@@ -161,7 +161,6 @@ pub fn gen_trace(
         .collect_vec()
 }
 
-// TODO(AlonH): Implement.
 impl ComponentTraceWriter<SimdBackend> for SimdWideFibComponent {
     fn write_interaction_trace(
         &self,
diff --git a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs b/crates/prover/src/examples/wide_fibonacci/trace_gen.rs
index 788fc842c..550aa0635 100644
--- a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs
+++ b/crates/prover/src/examples/wide_fibonacci/trace_gen.rs
@@ -1,5 +1,5 @@
 use itertools::Itertools;
-use num_traits::One;
+use num_traits::{One, Zero};
 
 use super::component::Input;
 use crate::core::fields::m31::BaseField;
@@ -49,23 +49,29 @@ pub fn write_lookup_column(
         })
         .collect_vec();
 
-    (0..n_rows)
+    let denominators = (0..n_rows)
         .map(|i| {
-            let numerator = shifted_secure_combination(
-                &[natural_ordered_trace[0][i], natural_ordered_trace[1][i]],
-                alpha,
-                z,
-            );
-            let denominator = shifted_secure_combination(
+            shifted_secure_combination(
                 &[
                     natural_ordered_trace[n_columns - 2][i],
                     natural_ordered_trace[n_columns - 1][i],
                 ],
                 alpha,
                 z,
+            )
+        })
+        .collect_vec();
+    let mut denominator_inverses = vec![SecureField::zero(); denominators.len()];
+    SecureField::batch_inverse(&denominators, &mut denominator_inverses);
+
+    (0..n_rows)
+        .map(|i| {
+            let numerator = shifted_secure_combination(
+                &[natural_ordered_trace[0][i], natural_ordered_trace[1][i]],
+                alpha,
+                z,
             );
-            // TODO(AlonH): Use batch inversion.
-            let cell = (numerator / denominator) * prev_value;
+            let cell = (numerator * denominator_inverses[i]) * prev_value;
             prev_value = cell;
             cell
         })