refactor: replace small mult with div lookup

zkonduit · Oct 5, 2023 · 87a03b0 · 87a03b0
1 parent fdb17a9
commit 87a03b0
Showing 1 changed file with 24 additions and 1 deletion.
diff --git a/src/graph/utilities.rs b/src/graph/utilities.rs
@@ -709,7 +709,30 @@ pub fn new_op_from_onnx(
         }
         "Add" => SupportedOp::Linear(PolyOp::Add),
         "Sub" => SupportedOp::Linear(PolyOp::Sub),
-        "Mul" => SupportedOp::Linear(PolyOp::Mult),
+        "Mul" => {
+            let mut op = SupportedOp::Linear(PolyOp::Mult);
+
+            let const_idx = inputs
+                .iter()
+                .enumerate()
+                .filter(|(_, n)| n.is_constant())
+                .map(|(i, _)| i)
+                .collect::<Vec<_>>();
+
+            assert_eq!(const_idx.len(), 1);
+            let const_idx = const_idx[0];
+
+            if let Some(c) = inputs[const_idx].opkind().get_mutable_constant() {
+                if c.raw_values.len() == 1 && c.raw_values[const_idx] < 1. {
+                    inputs[const_idx].decrement_const();
+                    deleted_indices.push(const_idx);
+                    op = SupportedOp::Nonlinear(LookupOp::Div {
+                        denom: crate::circuit::utils::F32(c.raw_values[0]),
+                    })
+                }
+            }
+            op
+        }
         "Iff" => SupportedOp::Linear(PolyOp::Iff),
         "Less" => {
             if inputs.len() == 2 {