small recursion program speedup

Author: TomWambsgans

crates/rec_aggregation/utils.py

@@ -81,20 +81,32 @@ def poly_eq_extension(point, n: Const):
     return res + (2**n - 1) * DIM
 
 
-def eq_mle_extension(a, b, n):
+@inline
+def eq_mle_extension_to(a, b, dst, n):
     debug_assert(n < 33)
     debug_assert(0 < n)
+    match_range(n, range(1, 33), lambda i: poly_eq_ee(a, b, dst, i))
+    return
+
+
+def eq_mle_extension(a, b, n):
     res = Array(DIM)
-    match_range(n, range(1, 33), lambda i: poly_eq_ee(a, b, res, i))
+    eq_mle_extension_to(a, b, res, n)
     return res
 
 
 @inline
-def eq_mle_base_extension(a, b, n):
+def eq_mle_base_extension_to(a, b, dst, n):
     debug_assert(n < 33)
     debug_assert(0 < n)
+    match_range(n, range(1, 33), lambda i: poly_eq_be(a, b, dst, i))
+    return
+
+
+@inline
+def eq_mle_base_extension(a, b, n):
     res = Array(DIM)
-    match_range(n, range(1, 33), lambda i: poly_eq_be(a, b, res, i))
+    eq_mle_base_extension_to(a, b, res, n)
     return res
 
 

crates/rec_aggregation/whir.py

@@ -113,8 +113,7 @@ def whir_open(
     all_ood_recovered_evals = Array(num_oods[0] * DIM)
     for i in range(0, num_oods[0]):
         expanded_from_univariate = expand_from_univariate_ext(ood_points_commit + i * DIM, n_vars)
-        ood_rec = eq_mle_extension(expanded_from_univariate, folding_randomness_global, n_vars)
-        copy_5(ood_rec, all_ood_recovered_evals + i * DIM)
+        eq_mle_extension_to(expanded_from_univariate, folding_randomness_global, all_ood_recovered_evals + i * DIM, n_vars)
     s: Mut = Array(DIM)
     dot_product_ee_dynamic(
         all_ood_recovered_evals,
@@ -132,8 +131,7 @@ def whir_open(
         my_folding_randomness += folding_factors[i] * DIM
         for j in range(0, num_oods[i + 1]):
             expanded_from_univariate = expand_from_univariate_ext(all_ood_points[i] + j * DIM, n_vars_remaining)
-            ood_rec = eq_mle_extension(expanded_from_univariate, my_folding_randomness, n_vars_remaining)
-            copy_5(ood_rec, my_ood_recovered_evals + j * DIM)
+            eq_mle_extension_to(expanded_from_univariate, my_folding_randomness, my_ood_recovered_evals + j * DIM, n_vars_remaining)
         summed_ood = Array(DIM)
         dot_product_ee_dynamic(
             my_ood_recovered_evals,
@@ -146,8 +144,7 @@ def whir_open(
         circle_value_i = all_circle_values[i]
         for j in range(0, num_queries[i]):  # unroll ?
             expanded_from_univariate = expand_from_univariate_base(circle_value_i[j], n_vars_remaining)
-            temp = eq_mle_base_extension(expanded_from_univariate, my_folding_randomness, n_vars_remaining)
-            copy_5(temp, s6s + j * DIM)
+            eq_mle_base_extension_to(expanded_from_univariate, my_folding_randomness, s6s + j * DIM, n_vars_remaining)
         s7 = Array(DIM)
         dot_product_ee_dynamic(
             s6s,