Dimitris/requantize bmm ref

.gitignore

@@ -1 +1 @@
-/target
+/target

Cargo.toml

@@ -19,15 +19,15 @@ rayon = { version = "1.5", default-features = false }
 ark-bn254 = { version = "^0.4.0", default-features = false, features = [ "curve" ] }
 blake2 = { version = "0.10", default-features = false }
 serde_json = "1.0.108"
-ark-pcs-bench-templates = { git = "https://github.com/HungryCatsStudio/poly-commit", branch = "ligero-uni-and-ml-absorb", default-features = false }
+ark-pcs-bench-templates = { git = "https://github.com/HungryCatsStudio/poly-commit", rev = "dfdd8e8" }
 more-asserts = { version = "0.3.1", default-features = false }
-pyo3 = { version = "0.20.3", features = ["auto-initialize"] }
+pyo3 = { version = "0.21.1", features = ["auto-initialize"] }
 
 [patch.crates-io]
-ark-ff = { git = "https://github.com/arkworks-rs/algebra/" }
-ark-ec = { git = "https://github.com/arkworks-rs/algebra/" }
-ark-serialize = { git = "https://github.com/arkworks-rs/algebra/" }
-ark-poly = { git = "https://github.com/arkworks-rs/algebra/" }
-ark-poly-commit = { git = "https://github.com/HungryCatsStudio/poly-commit", branch = "ligero-uni-and-ml-absorb" }
-ark-crypto-primitives = { git = "https://github.com/arkworks-rs/crypto-primitives" }
-ark-bn254 = { git = "https://github.com/arkworks-rs/algebra/" }
+ark-ff = { git = "https://github.com/HungryCatsStudio/algebra", rev = "97c7a4f" }
+ark-ec = { git = "https://github.com/HungryCatsStudio/algebra", rev = "97c7a4f" }
+ark-serialize = { git = "https://github.com/HungryCatsStudio/algebra", rev = "97c7a4f" }
+ark-poly = { git = "https://github.com/HungryCatsStudio/algebra", rev = "97c7a4f" }
+ark-poly-commit = { git = "https://github.com/HungryCatsStudio/poly-commit", rev = "bedc753" }
+ark-crypto-primitives = { git = "https://github.com/HungryCatsStudio/crypto-primitives", rev = "9ef1753" }
+ark-bn254 = { git = "https://github.com/HungryCatsStudio/algebra", rev = "97c7a4f" }

README.md

@@ -18,13 +18,13 @@ where `<example_name>` is one of the following:
 
 In order to run any tests involving python code, such as compatibility tests with TF Lite, the feature `python` must be activated (which automatically enables `test-types`).
 
-## From `ndarray` to `QArray`
+## From `ndarray` to `Tensor`
 
-In order to save a `numpy` `ndarray` (python side) as a serialised JSON which can be directly read into a `QArray` of ours (Rust side),
+In order to save a `numpy` `ndarray` (python side) as a serialised JSON which can be directly read into a `Tensor` of ours (Rust side),
 - Convert the `ndarray` into an `OrderedDict` using our custom python function `tensor_to_dict` (available in several of the python notebooks)
 - Pass the resulting `OrderedDict` together with the destination path to `json.dump`.
 
-The saved JSON file can be deserialised over in Rust with `QArray::read(path: &str) -> QArray`. If instead of a single `OrderedDict`, a python list of `OrderedDict`s is passed to `json.dump`, the resulting file can be deserialised with `QArray::read_list(path: &str) -> Vec<QArray> `.
+The saved JSON file can be deserialised over in Rust with `Tensor::read(path: &str) -> Tensor`. If instead of a single `OrderedDict`, a python list of `OrderedDict`s is passed to `json.dump`, the resulting file can be deserialised with `Tensor::read_list(path: &str) -> Vec<Tensor> `.
 
 Cf. `exploring_tf_lite/training_two_layer_perceptron.ipynb` for example usage.
 

common/examples/common/lib.rs

@@ -1,21 +1,21 @@
-use hcs_common::{quantise_f32_u8_nne, Model, Poly, QArray};
+use hcs_common::{quantise_f32_u8_nne, Model, Poly, Tensor};
 
 use ark_crypto_primitives::sponge::{Absorb, CryptographicSponge};
 use ark_ff::PrimeField;
 use ark_poly_commit::PolynomialCommitment;
 
 // Auxiliary function
 fn unpadded_inference<F, S, PCS>(
-    raw_input: QArray<f32>,
-    model: &Model<i8, i32>,
+    raw_input: Tensor<f32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
-) -> QArray<u8>
+) -> Tensor<u8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let quantised_input: QArray<u8> = QArray::new(
+    let quantised_input: Tensor<u8> = Tensor::new(
         quantise_f32_u8_nne(raw_input.values(), qinfo.0, qinfo.1),
         raw_input.shape().clone(),
     );
@@ -31,40 +31,40 @@ where
 // If padded inference is left on the prover side, move this to the prover
 /* // Auxiliary function
 fn padded_inference<F, S, PCS>(
-    raw_input: QArray<f32>,
-    model: &Model<i8, i32>,
+    raw_input: Tensor<f32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
-) -> QArray<u8>
+) -> Tensor<u8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let quantised_input: QArray<u8> = QArray::new(
+    let quantised_input: Tensor<u8> = Tensor::new(
         quantise_f32_u8_nne(raw_input.values(), qinfo.0, qinfo.1),
         raw_input.shape().clone(),
     );
 
     let input_i8 = (quantised_input.cast::<i32>() - 128).cast::<i8>();
 
     let output_i8 =
-        <Model<i8, i32> as ProveModel<F, S, PCS, i8, i32>>::padded_evaluate(model, input_i8);
+        <Model<i8> as ProveModel<F, S, PCS, i8>>::padded_evaluate(model, input_i8);
 
     (output_i8.cast::<i32>() + 128).cast()
 } */
 
 pub fn run_unpadded<F, S, PCS>(
     input_path: &str,
     expected_output_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_input: QArray<f32> = QArray::read(input_path);
-    let expected_output: QArray<u8> = QArray::read(expected_output_path);
+    let raw_input: Tensor<f32> = Tensor::read(input_path);
+    let expected_output: Tensor<u8> = Tensor::read(expected_output_path);
 
     let output_u8 = unpadded_inference::<F, S, PCS>(raw_input, model, qinfo);
 
@@ -78,15 +78,15 @@ pub fn run_unpadded<F, S, PCS>(
 /* pub fn run_padded<F, S, PCS>(
     input_path: &str,
     expected_output_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_input: QArray<f32> = QArray::read(input_path);
-    let expected_output: QArray<u8> = QArray::read(expected_output_path);
+    let raw_input: Tensor<f32> = Tensor::read(input_path);
+    let expected_output: Tensor<u8> = Tensor::read(expected_output_path);
 
     let output_u8 = padded_inference::<F, S, PCS>(raw_input, model, qinfo);
 
@@ -98,15 +98,15 @@ pub fn run_unpadded<F, S, PCS>(
 pub fn multi_run_unpadded<F, S, PCS>(
     inputs_path: &str,
     expected_outputs_path: &str,
-    model: &Model<i8, i32>,
+    model: &Model<i8>,
     qinfo: (f32, u8),
 ) where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_inputs: Vec<QArray<f32>> = QArray::read_list(inputs_path);
-    let expected_outputs: Vec<QArray<u8>> = QArray::read_list(expected_outputs_path);
+    let raw_inputs: Vec<Tensor<f32>> = Tensor::read_list(inputs_path);
+    let expected_outputs: Vec<Tensor<u8>> = Tensor::read_list(expected_outputs_path);
 
     for (raw_input, expected_output) in raw_inputs.into_iter().zip(expected_outputs.into_iter()) {
         assert_eq!(
@@ -131,8 +131,8 @@ pub fn multi_run_padded<F, S, PCS>(
     S: CryptographicSponge,
     PCS: PolynomialCommitment<F, Poly<F>, S>,
 {
-    let raw_inputs: Vec<QArray<f32>> = QArray::read_list(inputs_path);
-    let expected_outputs: Vec<QArray<u8>> = QArray::read_list(expected_outputs_path);
+    let raw_inputs: Vec<Tensor<f32>> = Tensor::read_list(inputs_path);
+    let expected_outputs: Vec<Tensor<u8>> = Tensor::read_list(expected_outputs_path);
 
     for (raw_input, expected_output) in raw_inputs.into_iter().zip(expected_outputs.into_iter()) {
         assert_eq!(

common/examples/simple_perceptron_mnist/main.rs

@@ -1,6 +1,6 @@
 use hcs_common::{
     simple_perceptron_mnist::{build_simple_perceptron_mnist, parameters::*},
-    Ligero,
+    BMMRequantizationStrategy, Ligero,
 };
 
 use ark_bn254::Fr;
@@ -20,7 +20,9 @@ macro_rules! PATH {
 }
 
 fn main() {
-    let simple_perceptron = build_simple_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>();
+    let simple_perceptron = build_simple_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>(
+        BMMRequantizationStrategy::Floating,
+    );
 
     // Right now this can't be QInfo because the latter is always a pair
     // (f32, i8), which indeed matches in-model quantisation, but not

common/examples/two_layer_perceptron_mnist/main.rs

@@ -1,6 +1,6 @@
 use hcs_common::{
     two_layer_perceptron_mnist::{build_two_layer_perceptron_mnist, parameters::*},
-    Ligero,
+    BMMRequantizationStrategy, Ligero,
 };
 
 use ark_bn254::Fr;
@@ -20,8 +20,9 @@ macro_rules! PATH {
 }
 
 fn main() {
-    let two_layer_perceptron =
-        build_two_layer_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>();
+    let two_layer_perceptron = build_two_layer_perceptron_mnist::<Fr, PoseidonSponge<Fr>, Ligero<Fr>>(
+        BMMRequantizationStrategy::Floating,
+    );
 
     // Right now this can't be QInfo because the latter is always a pair
     // (f32, i8), which indeed matches in-model quantisation, but not

common/src/compatibility/example_models/simple_perceptron_mnist/mod.rs

@@ -1,4 +1,10 @@
-use crate::{BMMNode, Model, Node, Poly, QArray, RequantiseBMMNode, ReshapeNode};
+use crate::{
+    model::nodes::{
+        requantize_bmm_ref::RequantizeBMMRefNode, requantize_bmm_single::RequantizeBMMSingleNode,
+    },
+    quantization::BMMRequantizationStrategy,
+    BMMNode, Model, Node, Poly, RequantizeBMMFloatNode, ReshapeNode, Tensor,
+};
 
 use ark_crypto_primitives::sponge::{Absorb, CryptographicSponge};
 use ark_ff::PrimeField;
@@ -22,7 +28,9 @@ macro_rules! PATH {
 }
 
 // TODO this is incorrect now that we have switched to logs
-pub fn build_simple_perceptron_mnist<F, S, PCS>() -> Model<i8, i32>
+pub fn build_simple_perceptron_mnist<F, S, PCS>(
+    req_strategy: BMMRequantizationStrategy,
+) -> Model<i8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
@@ -32,20 +40,25 @@ where
 
     let reshape: ReshapeNode = ReshapeNode::new(INPUT_DIMS.to_vec(), vec![flat_dim]);
 
-    let w_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights.json"));
-    let b_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias.json"));
+    let w_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights.json"));
+    let b_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias.json"));
 
-    let bmm: BMMNode<i8, i32> = BMMNode::new(w_array, b_array, Z_I);
+    let bmm: BMMNode<i8> = BMMNode::new(w_array, b_array, Z_I);
 
-    let req_bmm: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(OUTPUT_DIM, S_I, Z_I, S_W, Z_W, S_O, Z_O);
+    let req_bmm = match req_strategy {
+        BMMRequantizationStrategy::Floating => Node::RequantizeBMMFloat(
+            RequantizeBMMFloatNode::new(OUTPUT_DIM, S_I, Z_I, S_W, Z_W, S_O, Z_O),
+        ),
+        BMMRequantizationStrategy::Reference => {
+            Node::RequantizeBMMRef(RequantizeBMMRefNode::new(OUTPUT_DIM, S_I, S_W, S_O, Z_O))
+        }
+        BMMRequantizationStrategy::SingleRound => {
+            Node::RequantizeBMMSingle(RequantizeBMMSingleNode::new(OUTPUT_DIM, S_I, S_W, S_O, Z_O))
+        }
+    };
 
     Model::new(
         INPUT_DIMS.to_vec(),
-        vec![
-            Node::Reshape(reshape),
-            Node::BMM(bmm),
-            Node::RequantiseBMM(req_bmm),
-        ],
+        vec![Node::Reshape(reshape), Node::BMM(bmm), req_bmm],
     )
 }

common/src/compatibility/example_models/two_layer_perceptron_mnist/mod.rs

@@ -5,7 +5,13 @@ use ark_poly_commit::PolynomialCommitment;
 pub mod parameters;
 use parameters::*;
 
-use crate::{BMMNode, Model, Node, Poly, QArray, ReLUNode, RequantiseBMMNode, ReshapeNode};
+use crate::{
+    model::nodes::{
+        requantize_bmm_ref::RequantizeBMMRefNode, requantize_bmm_single::RequantizeBMMSingleNode,
+    },
+    quantization::BMMRequantizationStrategy,
+    BMMNode, Model, Node, Poly, ReLUNode, RequantizeBMMFloatNode, ReshapeNode, Tensor,
+};
 
 pub const INPUT_DIMS: &[usize] = &[28, 28];
 pub const INTER_DIM: usize = 28;
@@ -22,7 +28,9 @@ macro_rules! PATH {
     };
 }
 
-pub fn build_two_layer_perceptron_mnist<F, S, PCS>() -> Model<i8, i32>
+pub fn build_two_layer_perceptron_mnist<F, S, PCS>(
+    req_strategy: BMMRequantizationStrategy,
+) -> Model<i8>
 where
     F: PrimeField + Absorb,
     S: CryptographicSponge,
@@ -32,32 +40,50 @@ where
 
     let reshape: ReshapeNode = ReshapeNode::new(INPUT_DIMS.to_vec(), vec![flat_dim]);
 
-    let w1_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights_1.json"));
-    let b1_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias_1.json"));
-    let w2_array: QArray<i8> = QArray::read(&format!(PATH!(), "weights_2.json"));
-    let b2_array: QArray<i32> = QArray::read(&format!(PATH!(), "bias_2.json"));
+    let w1_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights_1.json"));
+    let b1_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias_1.json"));
+    let w2_array: Tensor<i8> = Tensor::read(&format!(PATH!(), "weights_2.json"));
+    let b2_array: Tensor<i32> = Tensor::read(&format!(PATH!(), "bias_2.json"));
 
-    let bmm_1: BMMNode<i8, i32> = BMMNode::new(w1_array, b1_array, Z_1_I);
+    let bmm_1: BMMNode<i8> = BMMNode::new(w1_array, b1_array, Z_1_I);
 
-    let req_bmm_1: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(INTER_DIM, S_1_I, Z_1_I, S_1_W, Z_1_W, S_1_O, Z_1_O);
+    let req_bmm_1 = match req_strategy {
+        BMMRequantizationStrategy::Floating => Node::RequantizeBMMFloat(
+            RequantizeBMMFloatNode::new(INTER_DIM, S_1_I, Z_1_I, S_1_W, Z_1_W, S_1_O, Z_1_O),
+        ),
+        BMMRequantizationStrategy::Reference => Node::RequantizeBMMRef(RequantizeBMMRefNode::new(
+            INTER_DIM, S_1_I, S_1_W, S_1_O, Z_1_O,
+        )),
+        BMMRequantizationStrategy::SingleRound => Node::RequantizeBMMSingle(
+            RequantizeBMMSingleNode::new(INTER_DIM, S_1_I, S_1_W, S_1_O, Z_1_O),
+        ),
+    };
 
     let relu: ReLUNode<i8> = ReLUNode::new(28, Z_1_O);
 
-    let bmm_2: BMMNode<i8, i32> = BMMNode::new(w2_array, b2_array, Z_2_I);
+    let bmm_2: BMMNode<i8> = BMMNode::new(w2_array, b2_array, Z_2_I);
 
-    let req_bmm_2: RequantiseBMMNode<i8> =
-        RequantiseBMMNode::new(OUTPUT_DIM, S_2_I, Z_2_I, S_2_W, Z_2_W, S_2_O, Z_2_O);
+    let req_bmm_2 = match req_strategy {
+        BMMRequantizationStrategy::Floating => Node::RequantizeBMMFloat(
+            RequantizeBMMFloatNode::new(OUTPUT_DIM, S_2_I, Z_2_I, S_2_W, Z_2_W, S_2_O, Z_2_O),
+        ),
+        BMMRequantizationStrategy::Reference => Node::RequantizeBMMRef(RequantizeBMMRefNode::new(
+            OUTPUT_DIM, S_2_I, S_2_W, S_2_O, Z_2_O,
+        )),
+        BMMRequantizationStrategy::SingleRound => Node::RequantizeBMMSingle(
+            RequantizeBMMSingleNode::new(OUTPUT_DIM, S_2_I, S_2_W, S_2_O, Z_2_O),
+        ),
+    };
 
     Model::new(
         INPUT_DIMS.to_vec(),
         vec![
             Node::Reshape(reshape),
             Node::BMM(bmm_1),
-            Node::RequantiseBMM(req_bmm_1),
+            req_bmm_1,
             Node::ReLU(relu),
             Node::BMM(bmm_2),
-            Node::RequantiseBMM(req_bmm_2),
+            req_bmm_2,
         ],
     )
 }