Merge branch 'perf_xmss'

Author: TomWambsgans

Cargo.lock

@@ -100,6 +100,19 @@ packed_pcs.workspace = true
 p3-air.workspace = true
 multilinear-toolkit.workspace = true
 
+# [patch."https://github.com/TomWambsgans/Plonky3.git"]
+# p3-koala-bear = { path = "../zk/Plonky3/koala-bear" }
+# p3-field = { path = "../zk/Plonky3/field" }
+# p3-poseidon2 = { path = "../zk/Plonky3/poseidon2" }
+# p3-matrix = { path = "../zk/Plonky3/matrix" }
+# p3-symmetric = { path = "../zk/Plonky3/symmetric" }
+# p3-air = { path = "../zk/Plonky3/air" }
+# p3-uni-stark = { path = "../zk/Plonky3/uni-stark" }
+# p3-poseidon2-air = { path = "../zk/Plonky3/poseidon2-air" }
+# p3-dft = { path = "../zk/Plonky3/dft" }
+# p3-challenger = { path = "../zk/Plonky3/challenger" }
+# p3-util = { path = "../zk/Plonky3/util" }
+
 [dev-dependencies]
 criterion = { version = "0.7", default-features = false, features = ["cargo_bench_support"] }
 rec_aggregation.workspace = true

Cargo.toml

@@ -12,7 +12,8 @@
 - avoid field embedding in the initial sumcheck of logup*, when table / values are in base field
 - opti logup* GKR when the indexes are not a power of 2 (which is the case in the execution table)
 - incremental merkle paths in whir-p3
-- Experiment to increase degree, and reduce commitments, in Poseidon arithmetization
+- Experiment to increase degree, and reduce commitments, in Poseidon arithmetization. 
+  Result: degree 9 is better than 3. TODO: degree 5, or 6 ? Also, the current degre 9 implem may not be perfectly optimal?
 - Avoid embedding overhead on the flag, len, and index columns in the AIR table for dot products
 - Batched logup*: when computing the eq() factor we can opti if the points contain boolean factor
 - Lev's trick to skip some low-level modular reduction
@@ -26,6 +27,7 @@
 - the interpreter of leanISA (+ witness generation) can be partially parallelized when there are some independent loops
 - (1 - x).r1 + x.r2 = x.(r2 - r1) + r1 TODO this opti is not everywhere currently + TODO generalize this with the univaraite skip
 - opti compute_eval_eq when scalar = ONE
+- Dmitry's range check, bonus: we can spare 2 memory cells if the value being range check is small (using the zeros present by conventio on the public memory)
 
 About "the packed pcs" (similar to SP1 Jagged PCS, slightly less efficient, but simpler (no sumchecks)):
 - The best strategy is probably to pack as much as possible (the cost increasing the density = additional inner evaluations), if we can fit below a power of 2 - epsilon  (epsilon = 20% for instance, tbd), if the sum of the non zero data is just above a power of 2, no packed technique, even the best, can help us, so we should spread aniway (to reduce the pressure of inner evaluations)

TODO.md

@@ -116,6 +116,15 @@ pub enum SimpleLine {
         to_decompose: Vec<SimpleExpr>,
         label: ConstMallocLabel,
     },
+    /// each field element x is decomposed to: (a0, a1, a2, ..., a11, b) where:
+    /// x = a0 + a1.4 + a2.4^2 + a3.4^3 + ... + a11.4^11 + b.2^24
+    /// and ai < 4, b < 2^7 - 1
+    /// The decomposition is unique, and always exists (except for x = -1)
+    DecomposeCustom {
+        var: Var, // a pointer to 13 * len(to_decompose) field elements
+        to_decompose: Vec<SimpleExpr>,
+        label: ConstMallocLabel,
+    },
     CounterHint {
         var: Var,
     },
@@ -653,6 +662,23 @@ fn simplify_lines(
                     label,
                 });
             }
+            Line::DecomposeCustom { var, to_decompose } => {
+                assert!(!const_malloc.forbidden_vars.contains(var), "TODO");
+                let simplified_to_decompose = to_decompose
+                    .iter()
+                    .map(|expr| {
+                        simplify_expr(expr, &mut res, counters, array_manager, const_malloc)
+                    })
+                    .collect::<Vec<_>>();
+                let label = const_malloc.counter;
+                const_malloc.counter += 1;
+                const_malloc.map.insert(var.clone(), label);
+                res.push(SimpleLine::DecomposeCustom {
+                    var: var.clone(),
+                    to_decompose: simplified_to_decompose,
+                    label,
+                });
+            }
             Line::CounterHint { var } => {
                 res.push(SimpleLine::CounterHint { var: var.clone() });
             }
@@ -837,12 +863,14 @@ pub fn find_variable_usage(lines: &[Line]) -> (BTreeSet<Var>, BTreeSet<Var>) {
                     on_new_expr(var, &internal_vars, &mut external_vars);
                 }
             }
-            Line::DecomposeBits { var, to_decompose } => {
+            Line::DecomposeBits { var, to_decompose }
+            | Line::DecomposeCustom { var, to_decompose } => {
                 for expr in to_decompose {
                     on_new_expr(expr, &internal_vars, &mut external_vars);
                 }
                 internal_vars.insert(var.clone());
             }
+
             Line::CounterHint { var } => {
                 internal_vars.insert(var.clone());
             }
@@ -1002,7 +1030,8 @@ pub fn inline_lines(
                     inline_expr(var, args, inlining_count);
                 }
             }
-            Line::DecomposeBits { var, to_decompose } => {
+            Line::DecomposeBits { var, to_decompose }
+            | Line::DecomposeCustom { var, to_decompose } => {
                 for expr in to_decompose {
                     inline_expr(expr, args, inlining_count);
                 }
@@ -1485,7 +1514,8 @@ fn replace_vars_for_unroll(
                     internal_vars,
                 );
             }
-            Line::DecomposeBits { var, to_decompose } => {
+            Line::DecomposeBits { var, to_decompose }
+            | Line::DecomposeCustom { var, to_decompose } => {
                 assert!(var != iterator, "Weird");
                 *var = format!("@unrolled_{unroll_index}_{iterator_value}_{var}");
                 for expr in to_decompose {
@@ -1886,6 +1916,7 @@ fn get_function_called(lines: &[Line], function_called: &mut Vec<String>) {
             | Line::Precompile { .. }
             | Line::Print { .. }
             | Line::DecomposeBits { .. }
+            | Line::DecomposeCustom { .. }
             | Line::CounterHint { .. }
             | Line::MAlloc { .. }
             | Line::Panic
@@ -1980,7 +2011,8 @@ fn replace_vars_by_const_in_lines(lines: &mut [Line], map: &BTreeMap<Var, F>) {
                     replace_vars_by_const_in_expr(var, map);
                 }
             }
-            Line::DecomposeBits { var, to_decompose } => {
+            Line::DecomposeBits { var, to_decompose }
+            | Line::DecomposeCustom { var, to_decompose } => {
                 assert!(!map.contains_key(var), "Variable {var} is a constant");
                 for expr in to_decompose {
                     replace_vars_by_const_in_expr(expr, map);
@@ -2063,6 +2095,21 @@ impl SimpleLine {
                         .join(", ")
                 )
             }
+            Self::DecomposeCustom {
+                var: result,
+                to_decompose,
+                label: _,
+            } => {
+                format!(
+                    "{} = decompose_custom({})",
+                    result,
+                    to_decompose
+                        .iter()
+                        .map(|expr| format!("{expr}"))
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                )
+            }
             Self::CounterHint { var: result } => {
                 format!("{result} = counter_hint()")
             }