From f3ebd687ab60557209c3c92ae6b7a141385a8446 Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 15:09:56 -0300 Subject: [PATCH 1/9] Pilopt: Optimize associative arithmetic operations --- pilopt/src/lib.rs | 83 +++++++++++++++++++++++++++++++++++++++ pilopt/tests/optimizer.rs | 29 ++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index 197f071ff4..20f3f9dc85 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -355,6 +355,14 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { return; } } + + if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) = e { + if let Some(simplified) = try_simplify_associative_operation(left, right, *op) { + *e = simplified; + return; + } + } + match e { AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, @@ -427,6 +435,81 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } } +fn try_simplify_associative_operation( + left: &Box>, + right: &Box>, + op: AlgebraicBinaryOperator, +) -> Option> { + // Find binary operation and other expression, handling both orderings: + // - (X ± C1) ± Other + // - Other ± (X ± C1) + let (bin_op, other_expr) = match (left.as_ref(), right.as_ref()) { + (AlgebraicExpression::BinaryOperation(bin), other) => (bin, other), + (other, AlgebraicExpression::BinaryOperation(bin)) => (bin, other), + _ => return None, + }; + + // Extract variable and constant from binary operation, handling both orderings: + // - (X ± C1) -> (X, C1) + // - (C1 ± X) -> (X, C1) + let AlgebraicBinaryOperation { + left: x1, + right: x2, + op: inner_op, + } = bin_op; + let (x, c1_val) = if let AlgebraicExpression::Number(val) = x1.as_ref() { + (x2, val) + } else if let AlgebraicExpression::Number(val) = x2.as_ref() { + (x1, val) + } else { + return None; + }; + + match other_expr { + // Case 1: Combining with a constant + // (X ± C1) ± C2 -> X ± (C1 ± C2) + // Handles all combinations of + and - for both operations + AlgebraicExpression::Number(c2) => { + let result = match inner_op { + AlgebraicBinaryOperator::Add => *c1_val + *c2, + AlgebraicBinaryOperator::Sub => *c1_val - *c2, + _ => return None, + }; + Some(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: x.clone(), + op, + right: Box::new(AlgebraicExpression::Number(result)), + }, + )) + } + + // Case 2: Combining with a binary operation + // (X ± C1) ± Y -> (X ± Y) ± C1 + // Preserves the sign of C1 based on the inner operation + y => { + let result = match inner_op { + AlgebraicBinaryOperator::Add => *c1_val, + AlgebraicBinaryOperator::Sub => -*c1_val, + _ => return None, + }; + Some(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: Box::new(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: x.clone(), + op, + right: Box::new(y.clone()), + }, + )), + op, + right: Box::new(AlgebraicExpression::Number(result)), + }, + )) + } + } +} + /// Extracts columns from lookups that are matched against constants and turns /// them into polynomial identities. fn extract_constant_lookups(pil_file: &mut Analyzed) { diff --git a/pilopt/tests/optimizer.rs b/pilopt/tests/optimizer.rs index 6a39eb24db..d6a0c30fde 100644 --- a/pilopt/tests/optimizer.rs +++ b/pilopt/tests/optimizer.rs @@ -465,3 +465,32 @@ fn equal_constrained_transitive() { let optimized = optimize(analyze_string::(input).unwrap()).to_string(); assert_eq!(optimized, expectation); } +#[test] +fn simplify_associative_operations() { + let input = r#"namespace N(150); + col witness x; + col witness y; + col witness z; + col fixed c1 = [1]*; + col fixed c2 = [2]*; + + (x + c2) + c1 = y; + (x - c2) + c1 = y; + + ((x + 3) - y) - 9 = z; + ((-x + 3) + y) + 9 = z; + "#; + + let expectation = r#"namespace N(150); + col witness x; + col witness y; + col witness z; + N::x - N::y = 3; + N::x - N::y = 1; + N::x - N::y - N::z = 6; + -N::x + N::y - N::z = 12; +"#; + + let optimized = optimize(analyze_string::(input).unwrap()).to_string(); + assert_eq!(optimized, expectation); +} From 01d2186a55a1631c4f8d93c57e72110f12c57191 Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 15:18:03 -0300 Subject: [PATCH 2/9] Clippy --- pilopt/src/lib.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index 20f3f9dc85..70c972200c 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -436,14 +436,14 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } fn try_simplify_associative_operation( - left: &Box>, - right: &Box>, + left: &AlgebraicExpression, + right: &AlgebraicExpression, op: AlgebraicBinaryOperator, ) -> Option> { // Find binary operation and other expression, handling both orderings: // - (X ± C1) ± Other // - Other ± (X ± C1) - let (bin_op, other_expr) = match (left.as_ref(), right.as_ref()) { + let (bin_op, other_expr) = match (left, right) { (AlgebraicExpression::BinaryOperation(bin), other) => (bin, other), (other, AlgebraicExpression::BinaryOperation(bin)) => (bin, other), _ => return None, @@ -468,7 +468,6 @@ fn try_simplify_associative_operation( match other_expr { // Case 1: Combining with a constant // (X ± C1) ± C2 -> X ± (C1 ± C2) - // Handles all combinations of + and - for both operations AlgebraicExpression::Number(c2) => { let result = match inner_op { AlgebraicBinaryOperator::Add => *c1_val + *c2, @@ -486,7 +485,6 @@ fn try_simplify_associative_operation( // Case 2: Combining with a binary operation // (X ± C1) ± Y -> (X ± Y) ± C1 - // Preserves the sign of C1 based on the inner operation y => { let result = match inner_op { AlgebraicBinaryOperator::Add => *c1_val, From 2d48c8b17eab86b2c9f2421859955d106c359e7f Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 17:15:52 -0300 Subject: [PATCH 3/9] Only where opt = Add --- pilopt/src/lib.rs | 75 +++++++++++++++++++-------------------- pilopt/tests/optimizer.rs | 8 ++--- 2 files changed, 40 insertions(+), 43 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index 70c972200c..b5306d7223 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -436,27 +436,35 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } fn try_simplify_associative_operation( - left: &AlgebraicExpression, - right: &AlgebraicExpression, + left: &Box>, + right: &Box>, op: AlgebraicBinaryOperator, ) -> Option> { + if op != AlgebraicBinaryOperator::Add { + return None; + } + // Find binary operation and other expression, handling both orderings: - // - (X ± C1) ± Other - // - Other ± (X ± C1) - let (bin_op, other_expr) = match (left, right) { + // (X + C1) + Other + // Other + (X + C1) + let (bin_op, other_expr) = match (left.as_ref(), right.as_ref()) { (AlgebraicExpression::BinaryOperation(bin), other) => (bin, other), (other, AlgebraicExpression::BinaryOperation(bin)) => (bin, other), _ => return None, }; - // Extract variable and constant from binary operation, handling both orderings: - // - (X ± C1) -> (X, C1) - // - (C1 ± X) -> (X, C1) let AlgebraicBinaryOperation { left: x1, right: x2, - op: inner_op, - } = bin_op; + op: inner_op @ AlgebraicBinaryOperator::Add, + } = bin_op + else { + return None; + }; + + // Extract variable and constant from binary operation, handling both orderings: + // (X + C1) -> (X, C1) + // (C1 + X) -> (X, C1) let (x, c1_val) = if let AlgebraicExpression::Number(val) = x1.as_ref() { (x2, val) } else if let AlgebraicExpression::Number(val) = x2.as_ref() { @@ -467,44 +475,33 @@ fn try_simplify_associative_operation( match other_expr { // Case 1: Combining with a constant - // (X ± C1) ± C2 -> X ± (C1 ± C2) + // (X + C1) + C2 -> X + (C1 + C2) AlgebraicExpression::Number(c2) => { - let result = match inner_op { - AlgebraicBinaryOperator::Add => *c1_val + *c2, - AlgebraicBinaryOperator::Sub => *c1_val - *c2, - _ => return None, - }; + let result = *c1_val + *c2; Some(AlgebraicExpression::BinaryOperation( AlgebraicBinaryOperation { left: x.clone(), - op, + op: AlgebraicBinaryOperator::Add, right: Box::new(AlgebraicExpression::Number(result)), }, )) } - // Case 2: Combining with a binary operation - // (X ± C1) ± Y -> (X ± Y) ± C1 - y => { - let result = match inner_op { - AlgebraicBinaryOperator::Add => *c1_val, - AlgebraicBinaryOperator::Sub => -*c1_val, - _ => return None, - }; - Some(AlgebraicExpression::BinaryOperation( - AlgebraicBinaryOperation { - left: Box::new(AlgebraicExpression::BinaryOperation( - AlgebraicBinaryOperation { - left: x.clone(), - op, - right: Box::new(y.clone()), - }, - )), - op, - right: Box::new(AlgebraicExpression::Number(result)), - }, - )) - } + // Case 2: Combining with any non-numeric expression + // (X + C1) + Y -> (X + Y) + C1 + y => Some(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: Box::new(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: x.clone(), + op: AlgebraicBinaryOperator::Add, + right: Box::new(y.clone()), + }, + )), + op: AlgebraicBinaryOperator::Add, + right: Box::new(AlgebraicExpression::Number(*c1_val)), + }, + )), } } diff --git a/pilopt/tests/optimizer.rs b/pilopt/tests/optimizer.rs index d6a0c30fde..e86311c749 100644 --- a/pilopt/tests/optimizer.rs +++ b/pilopt/tests/optimizer.rs @@ -485,10 +485,10 @@ fn simplify_associative_operations() { col witness x; col witness y; col witness z; - N::x - N::y = 3; - N::x - N::y = 1; - N::x - N::y - N::z = 6; - -N::x + N::y - N::z = 12; + N::x + 3 = N::y; + N::x - 2 + 1 = N::y; + N::x + 3 - N::y - 9 = N::z; + -N::x + N::y + 12 = N::z; "#; let optimized = optimize(analyze_string::(input).unwrap()).to_string(); From 878fb4dc4b407d4640fdf66f8fc0caac05b83a42 Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 17:24:41 -0300 Subject: [PATCH 4/9] Clippy --- pilopt/src/lib.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index b5306d7223..a19a044f76 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -436,8 +436,8 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } fn try_simplify_associative_operation( - left: &Box>, - right: &Box>, + left: &AlgebraicExpression, + right: &AlgebraicExpression, op: AlgebraicBinaryOperator, ) -> Option> { if op != AlgebraicBinaryOperator::Add { @@ -447,7 +447,7 @@ fn try_simplify_associative_operation( // Find binary operation and other expression, handling both orderings: // (X + C1) + Other // Other + (X + C1) - let (bin_op, other_expr) = match (left.as_ref(), right.as_ref()) { + let (bin_op, other_expr) = match (left, right) { (AlgebraicExpression::BinaryOperation(bin), other) => (bin, other), (other, AlgebraicExpression::BinaryOperation(bin)) => (bin, other), _ => return None, @@ -456,7 +456,7 @@ fn try_simplify_associative_operation( let AlgebraicBinaryOperation { left: x1, right: x2, - op: inner_op @ AlgebraicBinaryOperator::Add, + op: _inner_op @ AlgebraicBinaryOperator::Add, } = bin_op else { return None; From d3cff3fd346eb3b0ee0db150f565dcd9d671fdae Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 17:55:34 -0300 Subject: [PATCH 5/9] merge binary op matching and destructuring --- pilopt/src/lib.rs | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index a19a044f76..628e708efa 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -447,21 +447,26 @@ fn try_simplify_associative_operation( // Find binary operation and other expression, handling both orderings: // (X + C1) + Other // Other + (X + C1) - let (bin_op, other_expr) = match (left, right) { - (AlgebraicExpression::BinaryOperation(bin), other) => (bin, other), - (other, AlgebraicExpression::BinaryOperation(bin)) => (bin, other), + let (x1, x2, other_expr) = match (left, right) { + ( + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { + left: x1, + right: x2, + op: AlgebraicBinaryOperator::Add, + }), + other, + ) => (x1, x2, other), + ( + other, + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { + left: x1, + right: x2, + op: AlgebraicBinaryOperator::Add, + }), + ) => (x1, x2, other), _ => return None, }; - let AlgebraicBinaryOperation { - left: x1, - right: x2, - op: _inner_op @ AlgebraicBinaryOperator::Add, - } = bin_op - else { - return None; - }; - // Extract variable and constant from binary operation, handling both orderings: // (X + C1) -> (X, C1) // (C1 + X) -> (X, C1) From 2650c875a226bcab4013d56d00d0e6104e439f0e Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 30 Jan 2025 19:50:28 -0300 Subject: [PATCH 6/9] Test: more cases --- pilopt/tests/optimizer.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pilopt/tests/optimizer.rs b/pilopt/tests/optimizer.rs index e86311c749..75ea42246f 100644 --- a/pilopt/tests/optimizer.rs +++ b/pilopt/tests/optimizer.rs @@ -473,12 +473,16 @@ fn simplify_associative_operations() { col witness z; col fixed c1 = [1]*; col fixed c2 = [2]*; + col fixed c3 = [3]*; (x + c2) + c1 = y; + (c2 + x) + c3 = y; (x - c2) + c1 = y; ((x + 3) - y) - 9 = z; + (c3 + (x + 3)) - y = z; ((-x + 3) + y) + 9 = z; + ((-x + 3) + c3) + 12 = z; "#; let expectation = r#"namespace N(150); @@ -486,9 +490,12 @@ fn simplify_associative_operations() { col witness y; col witness z; N::x + 3 = N::y; + N::x + 5 = N::y; N::x - 2 + 1 = N::y; N::x + 3 - N::y - 9 = N::z; + N::x + 6 - N::y = N::z; -N::x + N::y + 12 = N::z; + -N::x + 18 = N::z; "#; let optimized = optimize(analyze_string::(input).unwrap()).to_string(); From b51e207c9dc8a5c39983936de7ee7a9d770a49f6 Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Fri, 31 Jan 2025 10:47:01 -0300 Subject: [PATCH 7/9] Docs & avoid clone --- pilopt/src/lib.rs | 41 +++++++++++++++++++++++------------------ 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index 628e708efa..d06b448ce7 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -435,9 +435,10 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } } +/// Uses associative properties to simplify expressions by regrouping constant terms. fn try_simplify_associative_operation( - left: &AlgebraicExpression, - right: &AlgebraicExpression, + left: &mut AlgebraicExpression, + right: &mut AlgebraicExpression, op: AlgebraicBinaryOperator, ) -> Option> { if op != AlgebraicBinaryOperator::Add { @@ -471,13 +472,14 @@ fn try_simplify_associative_operation( // (X + C1) -> (X, C1) // (C1 + X) -> (X, C1) let (x, c1_val) = if let AlgebraicExpression::Number(val) = x1.as_ref() { - (x2, val) + (x2.as_mut(), val) } else if let AlgebraicExpression::Number(val) = x2.as_ref() { - (x1, val) + (x1.as_mut(), val) } else { return None; }; + let x = std::mem::replace(x, AlgebraicExpression::Number(0.into())); match other_expr { // Case 1: Combining with a constant // (X + C1) + C2 -> X + (C1 + C2) @@ -485,7 +487,7 @@ fn try_simplify_associative_operation( let result = *c1_val + *c2; Some(AlgebraicExpression::BinaryOperation( AlgebraicBinaryOperation { - left: x.clone(), + left: Box::new(x), op: AlgebraicBinaryOperator::Add, right: Box::new(AlgebraicExpression::Number(result)), }, @@ -494,19 +496,22 @@ fn try_simplify_associative_operation( // Case 2: Combining with any non-numeric expression // (X + C1) + Y -> (X + Y) + C1 - y => Some(AlgebraicExpression::BinaryOperation( - AlgebraicBinaryOperation { - left: Box::new(AlgebraicExpression::BinaryOperation( - AlgebraicBinaryOperation { - left: x.clone(), - op: AlgebraicBinaryOperator::Add, - right: Box::new(y.clone()), - }, - )), - op: AlgebraicBinaryOperator::Add, - right: Box::new(AlgebraicExpression::Number(*c1_val)), - }, - )), + y => { + let y = std::mem::replace(y, AlgebraicExpression::Number(0.into())); + Some(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: Box::new(AlgebraicExpression::BinaryOperation( + AlgebraicBinaryOperation { + left: Box::new(x), + op: AlgebraicBinaryOperator::Add, + right: Box::new(y), + }, + )), + op: AlgebraicBinaryOperator::Add, + right: Box::new(AlgebraicExpression::Number(*c1_val)), + }, + )) + } } } From 269e3194f4cb5ae9e1d2aac5d3d68be6f67c4d4c Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Mon, 3 Feb 2025 11:58:07 -0300 Subject: [PATCH 8/9] Comments improved --- pilopt/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index d06b448ce7..34069647e4 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -446,8 +446,8 @@ fn try_simplify_associative_operation( } // Find binary operation and other expression, handling both orderings: - // (X + C1) + Other - // Other + (X + C1) + // (X1 + X2) + Other + // Other + (X1 + X2) let (x1, x2, other_expr) = match (left, right) { ( AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { @@ -469,8 +469,8 @@ fn try_simplify_associative_operation( }; // Extract variable and constant from binary operation, handling both orderings: - // (X + C1) -> (X, C1) - // (C1 + X) -> (X, C1) + // (X1 + C1) -> (X1, C1) if X2 is a constant + // (C1 + X2) -> (X2, C1) if X1 is a constant let (x, c1_val) = if let AlgebraicExpression::Number(val) = x1.as_ref() { (x2.as_mut(), val) } else if let AlgebraicExpression::Number(val) = x2.as_ref() { @@ -482,7 +482,7 @@ fn try_simplify_associative_operation( let x = std::mem::replace(x, AlgebraicExpression::Number(0.into())); match other_expr { // Case 1: Combining with a constant - // (X + C1) + C2 -> X + (C1 + C2) + // (X + C1) + Other -> X + (C1 + Other) AlgebraicExpression::Number(c2) => { let result = *c1_val + *c2; Some(AlgebraicExpression::BinaryOperation( From 9e7e71d4f2c7c58dcd6db941e99d2b8c140c7c47 Mon Sep 17 00:00:00 2001 From: Gaston Zanitti Date: Thu, 6 Feb 2025 16:35:01 -0300 Subject: [PATCH 9/9] BinaryOp & extra BinaryOp removed --- pilopt/src/lib.rs | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/pilopt/src/lib.rs b/pilopt/src/lib.rs index 34069647e4..69474eddcf 100644 --- a/pilopt/src/lib.rs +++ b/pilopt/src/lib.rs @@ -356,8 +356,8 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } } - if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) = e { - if let Some(simplified) = try_simplify_associative_operation(left, right, *op) { + if let AlgebraicExpression::BinaryOperation(binary_op) = e { + if let Some(simplified) = try_simplify_associative_operation(binary_op) { *e = simplified; return; } @@ -435,20 +435,17 @@ fn simplify_expression_single(e: &mut AlgebraicExpression) { } } -/// Uses associative properties to simplify expressions by regrouping constant terms. fn try_simplify_associative_operation( - left: &mut AlgebraicExpression, - right: &mut AlgebraicExpression, - op: AlgebraicBinaryOperator, + binary_op: &mut AlgebraicBinaryOperation, ) -> Option> { - if op != AlgebraicBinaryOperator::Add { + if binary_op.op != AlgebraicBinaryOperator::Add { return None; } // Find binary operation and other expression, handling both orderings: // (X1 + X2) + Other // Other + (X1 + X2) - let (x1, x2, other_expr) = match (left, right) { + let (x1, x2, other_expr) = match (&mut *binary_op.left, &mut *binary_op.right) { ( AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: x1, @@ -500,13 +497,7 @@ fn try_simplify_associative_operation( let y = std::mem::replace(y, AlgebraicExpression::Number(0.into())); Some(AlgebraicExpression::BinaryOperation( AlgebraicBinaryOperation { - left: Box::new(AlgebraicExpression::BinaryOperation( - AlgebraicBinaryOperation { - left: Box::new(x), - op: AlgebraicBinaryOperator::Add, - right: Box::new(y), - }, - )), + left: Box::new(x + y), op: AlgebraicBinaryOperator::Add, right: Box::new(AlgebraicExpression::Number(*c1_val)), }, @@ -514,7 +505,6 @@ fn try_simplify_associative_operation( } } } - /// Extracts columns from lookups that are matched against constants and turns /// them into polynomial identities. fn extract_constant_lookups(pil_file: &mut Analyzed) {