ciphercore_base/ops/
newton_inversion.rs

1//! Multiplicative inversion via [the Newton-Raphson method](https://en.wikipedia.org/wiki/Newton%27s_method#Multiplicative_inverses_of_numbers_and_power_series).
2use crate::custom_ops::CustomOperationBody;
3use crate::data_types::{Type, INT64, UINT64};
4use crate::errors::Result;
5use crate::graphs::{Context, Graph};
6
7use serde::{Deserialize, Serialize};
8
9use super::utils::{constant_scalar, inverse_initial_approximation, multiply_fixed_point};
10
11/// A structure that defines the custom operation NewtonInversion that computes an inversion of a number via [the Newton-Raphson method](https://en.wikipedia.org/wiki/Newton%27s_method#Multiplicative_inverses_of_numbers_and_power_series).
12///
13/// In particular, this operation computes an approximation of 2<sup>denominator_cap_2k</sup> / input.
14///
15/// Input must be of the scalar type UINT64 or INT64 and be in (0, 2<sup>denominator_cap_2k - 1</sup>) range.
16/// The input is also assumed to be small enough (less than 2<sup>32</sup>), otherwise integer overflows
17/// are possible, yielding incorrect results.
18///
19/// Optionally, an initial approximation for the Newton-Raphson method can be provided.
20/// In this case, the operation might be faster and of lower depth, however, it must be guaranteed that
21/// 2<sup>denominator_cap_2k - 1</sup> <= input * initial_approximation < 2<sup>denominator_cap_2k + 1</sup>.
22///
23/// # Custom operation arguments
24///
25/// - Node containing an unsigned or signed 64-bit array or scalar to invert
26/// - Negative values are currently unsupported as sign extraction is quite expensive
27/// - (optional) Node containing an array or scalar that serves as an initial approximation of the Newton-Raphson method
28///
29/// # Custom operation returns
30///
31/// New NewtonInversion node
32///
33/// # Example
34///
35/// ```
36/// # use ciphercore_base::graphs::create_context;
37/// # use ciphercore_base::data_types::{scalar_type, array_type, UINT64};
38/// # use ciphercore_base::custom_ops::{CustomOperation};
39/// # use ciphercore_base::ops::newton_inversion::NewtonInversion;
40/// let c = create_context().unwrap();
41/// let g = c.create_graph().unwrap();
42/// let t = array_type(vec![2, 3], UINT64);
43/// let n1 = g.input(t.clone()).unwrap();
44/// let guess_n = g.input(t.clone()).unwrap();
45/// let n2 = g.custom_op(CustomOperation::new(NewtonInversion {iterations: 10, denominator_cap_2k: 4}), vec![n1, guess_n]).unwrap();
46///
47// TODO: generalize to other types.
48#[derive(Debug, Serialize, Deserialize, Eq, PartialEq, Hash)]
49pub struct NewtonInversion {
50    /// Number of iterations of the Newton-Raphson method; rule of thumb is to set it to 1 + log(`denominator_cap_2k`)
51    pub iterations: u64,
52    /// Number of output bits that are approximated
53    pub denominator_cap_2k: u64,
54}
55
56#[typetag::serde]
57impl CustomOperationBody for NewtonInversion {
58    fn instantiate(&self, context: Context, arguments_types: Vec<Type>) -> Result<Graph> {
59        if arguments_types.len() != 1 && arguments_types.len() != 2 {
60            return Err(runtime_error!(
61                "Invalid number of arguments for NewtonDivision"
62            ));
63        }
64        let t = arguments_types[0].clone();
65        if !t.is_scalar() && !t.is_array() {
66            return Err(runtime_error!(
67                "Divisor in NewtonDivision must be a scalar or an array"
68            ));
69        }
70        let sc = t.get_scalar_type();
71        if sc != UINT64 && sc != INT64 {
72            return Err(runtime_error!(
73                "Divisor in NewtonDivision must consist of either INT64s or UINT64s"
74            ));
75        }
76        let has_initial_approximation = arguments_types.len() == 2;
77        if has_initial_approximation {
78            let divisor_t = arguments_types[1].clone();
79            if divisor_t != t {
80                return Err(runtime_error!(
81                    "Divisor and initial approximation must have the same type."
82                ));
83            }
84        }
85
86        let g_initial_approximation =
87            inverse_initial_approximation(&context, t.clone(), self.denominator_cap_2k)?;
88        let g = context.create_graph()?;
89        let divisor = g.input(t.clone())?;
90        let mut approximation = if has_initial_approximation {
91            g.input(t)?
92        } else if self.denominator_cap_2k == 0 {
93            g.ones(t)?
94        } else {
95            g.call(g_initial_approximation, vec![divisor.clone()])?
96        };
97        // Now, we do Newton approximation for computing 1 / x, where x = divisor / (2 ** cap).
98        // The formula for the Newton method is x_{i + 1} = x_i * (2 - d * x_i).
99        let two_power_cap_plus_one = constant_scalar(&g, 1 << (self.denominator_cap_2k + 1), sc)?;
100        for _ in 0..self.iterations {
101            let x = approximation;
102            let mult = two_power_cap_plus_one.subtract(x.multiply(divisor.clone())?)?;
103            approximation = multiply_fixed_point(mult, x, self.denominator_cap_2k)?;
104        }
105        approximation.set_as_output()?;
106        g.finalize()?;
107        Ok(g)
108    }
109
110    fn get_name(&self) -> String {
111        format!(
112            "NewtonDivision(iterations={}, cap=2**{})",
113            self.iterations, self.denominator_cap_2k
114        )
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    use crate::custom_ops::run_instantiation_pass;
123    use crate::custom_ops::CustomOperation;
124    use crate::data_types::{array_type, scalar_type, ScalarType};
125    use crate::data_values::Value;
126    use crate::evaluators::random_evaluate;
127    use crate::graphs::util::simple_context;
128    use crate::inline::inline_common::DepthOptimizationLevel;
129    use crate::inline::inline_ops::inline_operations;
130    use crate::inline::inline_ops::InlineConfig;
131    use crate::inline::inline_ops::InlineMode;
132    use crate::mpc::mpc_compiler::prepare_for_mpc_evaluation;
133    use crate::mpc::mpc_compiler::IOStatus;
134
135    fn scalar_division_helper(
136        divisor: u64,
137        initial_approximation: Option<u64>,
138        st: ScalarType,
139    ) -> Result<Value> {
140        let c = simple_context(|g| {
141            let i = g.input(scalar_type(st))?;
142            if let Some(approx) = initial_approximation {
143                let approx_const = constant_scalar(&g, approx, st)?;
144                g.custom_op(
145                    CustomOperation::new(NewtonInversion {
146                        iterations: 5,
147                        denominator_cap_2k: 10,
148                    }),
149                    vec![i, approx_const],
150                )
151            } else {
152                g.custom_op(
153                    CustomOperation::new(NewtonInversion {
154                        iterations: 5,
155                        denominator_cap_2k: 10,
156                    }),
157                    vec![i],
158                )
159            }
160        })?;
161        let mapped_c = run_instantiation_pass(c)?;
162        let result = random_evaluate(
163            mapped_c.get_context().get_main_graph()?,
164            vec![Value::from_scalar(divisor, st)?],
165        )?;
166        Ok(result)
167    }
168
169    fn array_division_helper(divisor: Vec<u64>, st: ScalarType) -> Result<Vec<u64>> {
170        let array_t = array_type(vec![divisor.len() as u64], st);
171        let c = simple_context(|g| {
172            let i = g.input(array_t.clone())?;
173            g.custom_op(
174                CustomOperation::new(NewtonInversion {
175                    iterations: 5,
176                    denominator_cap_2k: 10,
177                }),
178                vec![i],
179            )
180        })?;
181        let mapped_c = run_instantiation_pass(c)?;
182        let result = random_evaluate(
183            mapped_c.get_context().get_main_graph()?,
184            vec![Value::from_flattened_array(&divisor, st)?],
185        )?;
186        result.to_flattened_array_u64(array_t)
187    }
188
189    #[test]
190    fn test_newton_division_scalar() {
191        let div_v = vec![1, 2, 3, 123, 300, 500, 700];
192        for i in div_v.clone() {
193            assert!(
194                (scalar_division_helper(i, None, UINT64)
195                    .unwrap()
196                    .to_u64(UINT64)
197                    .unwrap() as i64
198                    - 1024 / i as i64)
199                    .abs()
200                    <= 1
201            );
202
203            assert!(
204                (scalar_division_helper(i, None, INT64)
205                    .unwrap()
206                    .to_i64(INT64)
207                    .unwrap() as i64
208                    - 1024 / i as i64)
209                    .abs()
210                    <= 1
211            );
212        }
213    }
214
215    #[test]
216    fn test_newton_division_array() {
217        let arr = vec![23, 32, 57, 71, 183, 555];
218        let div = array_division_helper(arr.clone(), UINT64).unwrap();
219        let i_div = array_division_helper(arr.clone(), INT64).unwrap();
220        for i in 0..arr.len() {
221            assert!((div[i] as i64 - 1024 / arr[i] as i64).abs() <= 1);
222            assert!((i_div[i] as i64 - 1024 / arr[i] as i64).abs() <= 1);
223        }
224    }
225
226    #[test]
227    fn test_newton_division_with_initial_guess() {
228        for i in vec![1, 2, 3, 123, 300, 500, 700] {
229            let mut initial_guess = 1;
230            while initial_guess * i * 2 < 1024 {
231                initial_guess *= 2;
232            }
233            assert!(
234                (scalar_division_helper(i, Some(initial_guess), UINT64)
235                    .unwrap()
236                    .to_u64(UINT64)
237                    .unwrap() as i64
238                    - 1024 / i as i64)
239                    .abs()
240                    <= 1
241            );
242            assert!(
243                (scalar_division_helper(i, Some(initial_guess), INT64)
244                    .unwrap()
245                    .to_i64(INT64)
246                    .unwrap() as i64
247                    - 1024 / i as i64)
248                    .abs()
249                    <= 1
250            );
251        }
252    }
253
254    #[test]
255    fn test_newton_inversion_compiles_end2end() -> Result<()> {
256        let c = simple_context(|g| {
257            let i = g.input(scalar_type(INT64))?;
258            g.custom_op(
259                CustomOperation::new(NewtonInversion {
260                    iterations: 5,
261                    denominator_cap_2k: 10,
262                }),
263                vec![i],
264            )
265        })?;
266        let inline_config = InlineConfig {
267            default_mode: InlineMode::DepthOptimized(DepthOptimizationLevel::Default),
268            ..Default::default()
269        };
270        let instantiated_context = run_instantiation_pass(c)?.get_context();
271        let inlined_context = inline_operations(instantiated_context, inline_config.clone())?;
272        let _unused = prepare_for_mpc_evaluation(
273            inlined_context,
274            vec![vec![IOStatus::Shared]],
275            vec![vec![]],
276            inline_config,
277        )?;
278        Ok(())
279    }
280}
ciphercore_base/ops/newton_inversion.rs

ciphercore_base/ops/
newton_inversion.rs
