Repository: ddkang/zkml Branch: main Commit: 43789582671f Files: 93 Total size: 13.7 MB Directory structure: gitextract_ixvshx2q/ ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── backwards/ │ ├── README.md │ └── backward.py ├── python/ │ ├── converter.py │ ├── input_converter.py │ └── training_converter.py ├── rustfmt.toml ├── src/ │ ├── bin/ │ │ ├── test_circuit.rs │ │ ├── time_circuit.rs │ │ ├── verify_circuit.rs │ │ └── verify_wav.rs │ ├── commitments/ │ │ ├── commit.rs │ │ ├── packer.rs │ │ └── poseidon_commit.rs │ ├── commitments.rs │ ├── gadgets/ │ │ ├── add_pairs.rs │ │ ├── adder.rs │ │ ├── bias_div_floor_relu6.rs │ │ ├── bias_div_round_relu6.rs │ │ ├── dot_prod.rs │ │ ├── gadget.rs │ │ ├── input_lookup.rs │ │ ├── max.rs │ │ ├── mul_pairs.rs │ │ ├── nonlinear/ │ │ │ ├── exp.rs │ │ │ ├── logistic.rs │ │ │ ├── non_linearity.rs │ │ │ ├── pow.rs │ │ │ ├── relu.rs │ │ │ ├── rsqrt.rs │ │ │ ├── sqrt.rs │ │ │ └── tanh.rs │ │ ├── nonlinear.rs │ │ ├── sqrt_big.rs │ │ ├── square.rs │ │ ├── squared_diff.rs │ │ ├── sub_pairs.rs │ │ ├── update.rs │ │ ├── var_div.rs │ │ ├── var_div_big.rs │ │ └── var_div_big3.rs │ ├── gadgets.rs │ ├── layers/ │ │ ├── arithmetic/ │ │ │ ├── add.rs │ │ │ ├── div_var.rs │ │ │ ├── mul.rs │ │ │ └── sub.rs │ │ ├── arithmetic.rs │ │ ├── averager.rs │ │ ├── avg_pool_2d.rs │ │ ├── batch_mat_mul.rs │ │ ├── conv2d.rs │ │ ├── dag.rs │ │ ├── div_fixed.rs │ │ ├── fully_connected.rs │ │ ├── layer.rs │ │ ├── logistic.rs │ │ ├── max_pool_2d.rs │ │ ├── mean.rs │ │ ├── noop.rs │ │ ├── pow.rs │ │ ├── rsqrt.rs │ │ ├── shape/ │ │ │ ├── broadcast.rs │ │ │ ├── concatenation.rs │ │ │ ├── mask_neg_inf.rs │ │ │ ├── pack.rs │ │ │ ├── pad.rs │ │ │ ├── permute.rs │ │ │ ├── reshape.rs │ │ │ ├── resize_nn.rs │ │ │ ├── rotate.rs │ │ │ ├── slice.rs │ │ │ ├── split.rs │ │ │ └── transpose.rs │ │ ├── shape.rs │ │ ├── softmax.rs │ │ ├── sqrt.rs │ │ ├── square.rs │ │ ├── squared_diff.rs │ │ ├── tanh.rs │ │ └── update.rs │ ├── layers.rs │ ├── lib.rs │ ├── model.rs │ ├── utils/ │ │ ├── helpers.rs │ │ ├── loader.rs │ │ ├── proving_ipa.rs │ │ └── proving_kzg.rs │ └── utils.rs └── testing/ └── circuits/ ├── last_two_layers.py └── v2_1.0_224.tflite ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ # Generated by Cargo # will have compiled files and executables /target/ # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html Cargo.lock # These are backup files generated by rustfmt **/*.rs.bk # VSCode .vscode **/.DS_Store *.swp # Proof stuff out.msgpack proof_size_check pkey vkey proof params_kzg params_ipa examples testing/data *.diff ================================================ FILE: Cargo.toml ================================================ [package] name = "zkml" version = "0.0.1" edition = "2021" description = "Zero-knowledge machine learning" license = "LICENSE" homepage = "https://github.com/ddkang/zkml" repository = "https://github.com/ddkang/zkml-public.git" readme = "README.md" exclude = [ "params", "params_kzg", "python", ] [profile.dev] opt-level = 3 [profile.test] opt-level = 3 [dependencies] bitvec = "1.0.1" halo2 = { git="https://github.com/privacy-scaling-explorations/halo2", package="halo2", rev="17e9765c199670534c0299c96128d0464a188d0b" } halo2_gadgets = { git="https://github.com/privacy-scaling-explorations/halo2", package="halo2_gadgets", rev="17e9765c199670534c0299c96128d0464a188d0b", features = ["circuit-params"] } halo2_proofs = { git="https://github.com/privacy-scaling-explorations/halo2", package="halo2_proofs", rev="17e9765c199670534c0299c96128d0464a188d0b", features = ["circuit-params"] } lazy_static = "1.4.0" ndarray = "0.15.6" num-bigint = "0.4.3" num-traits = "0.2.15" once_cell = "1.15.0" rand = "0.8.5" rmp-serde = "1.1.1" rounded-div = "0.1.2" serde = "1.0.152" serde_derive = "1.0.152" serde_json = "1.0.85" wav = "1.0.0" ================================================ FILE: LICENSE ================================================ Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ================================================ FILE: README.md ================================================ # zkml zkml is a framework for constructing proofs of ML model execution in ZK-SNARKs. Read our [blog post](https://medium.com/@danieldkang/trustless-verification-of-machine-learning-6f648fd8ba88) and [paper](https://arxiv.org/abs/2210.08674) for implementation details. zkml requires the nightly build of Rust: ``` rustup override set nightly ``` ## Quickstart Run the following commands: ```sh # Installs rust, skip if you already have rust installed curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh git clone https://github.com/ddkang/zkml.git cd zkml rustup override set nightly cargo build --release mkdir params_kzg mkdir params_ipa # This should take ~16s to run the first time # and ~8s to run the second time ./target/release/time_circuit examples/mnist/model.msgpack examples/mnist/inp.msgpack kzg ``` This will prove an MNIST circuit! It will require around 2GB of memory and take around 8 seconds to run. ## Converting your own model and data To convert your own model and data, you will need to convert the model and data to the format zkml expects. Currently, we accept TFLite models. We show an example below. 1. First `cd examples/mnist` 2. We've already created a model that achieves high accuracy on MNIST (`model.tflite`). You will need to create your own TFLite model. One way is to [convert a model from Keras](https://stackoverflow.com/questions/53256877/how-to-convert-kerash5-file-to-a-tflite-file). 3. You will need to convert the model: ```bash python ../../python/converter.py --model model.tflite --model_output converted_model.msgpack --config_output config.msgpack --scale_factor 512 --k 17 --num_cols 10 --num_randoms 1024 ``` There are several parameters that need to be changed depending on the model (`scale_factor`, `k`, `num_cols`, and `num_randoms`). 4. You will first need to serialize the model input to numpy's serialization format `npy`. We've written a small script to do this for the first test data point in MNIST: ```bash python data_to_npy.py ``` 5. You will then need to convert the input to the model: ```bash python ../../python/input_converter.py --model_config converted_model.msgpack --inputs 7.npy --output example_inp.msgpack ``` 6. Once you've converted the model and input, you can run the model as above! However, we generally recommend testing the model before proving (you will need to build zkml before running the next line): ```bash cd ../../ ./target/release/test_circuit examples/mnist/converted_model.msgpack examples/mnist/example_inp.msgpack ``` ## Contact us If you're interested in extending or using zkml, please contact us at `ddkang [at] g.illinois.edu`. ================================================ FILE: backwards/README.md ================================================ ### About Takes in a feed-forward TF model and outputs a new computational graph for back-propagation. ================================================ FILE: backwards/backward.py ================================================ # # A script for generating a backprop computational graph from forward # import argparse import ast from typing import Literal, Union import msgpack import numpy as np class CircuitConfig(): def __init__(self, starting_index): self.next_index = starting_index self.outp_to_grad = {} self.label_tensor_idx = None self.weights_update = None # Allocates an index for a gradient tensor and returns def new_gradient_tensor(self, tensor_idx): if tensor_idx in self.outp_to_grad: raise Exception("Tensor already allocated") self.outp_to_grad[tensor_idx] = self.next_index self.next_index += 1 return self.outp_to_grad[tensor_idx] # Allocates an index for a tensor def new_tensor(self): new_index = self.next_index self.next_index += 1 return new_index def new_label_tensor(self): if self.label_tensor_idx is not None: raise Exception("Label tensor already allocated") self.label_tensor_idx = self.next_index self.next_index += 1 return self.label_tensor_idx # Allocates an index for a gradient tensor and returns def gradient_tensor_idx(self, tensor_idx): return self.outp_to_grad[tensor_idx] # TODO: Put these in enums NO_ACTIVATION = 0 SAME = 0 VALID = 1 CONV2D = 0 CONV2D_DEPTHWISE = 1 class Conv2D(): def __init__(self, layer): params = layer['params'] self.padding = params[1] self.activation_type = params[2] self.stride_h = params[3] self.stride_w = params[4] def backward(self, layer, transcript, config): inputs_idx, inputs_shape = layer['inp_idxes'][0], layer['inp_shapes'][0] weights_idx, weights_shape = layer['inp_idxes'][1], layer['inp_shapes'][1] bias_idx, bias_shape = layer['inp_idxes'][2], layer['inp_shapes'][2] output_idx, output_shape = layer['out_idxes'][0], layer['out_shapes'][0] permuted_inputs_idx = config.new_tensor() permutation = [3, 1, 2, 0] permuted_inputs_shape = [inputs_shape[p] for p in permutation] inputs_permute_layer = { 'layer_type': 'Permute', 'params': permutation, 'inp_idxes': [inputs_idx], 'out_idxes': [permuted_inputs_idx], 'inp_shapes': [inputs_shape], 'out_shapes': [permuted_inputs_shape], 'mask': [], } transcript.append(inputs_permute_layer) permuted_outputs_idx = config.new_tensor() permuted_outputs_shape = [output_shape[p] for p in permutation] inputs_permute_layer = { 'layer_type': 'Permute', 'params': permutation, 'inp_idxes': [config.gradient_tensor_idx(output_idx)], 'out_idxes': [permuted_outputs_idx], 'inp_shapes': [output_shape], 'out_shapes': [permuted_outputs_shape], 'mask': [], } transcript.append(inputs_permute_layer) dw_idx, dw_shape = config.new_tensor(), weights_shape dw_conv = { 'layer_type': 'Conv2D', 'params': [CONV2D, VALID, NO_ACTIVATION, self.stride_h, self.stride_w], 'inp_idxes': [permuted_inputs_idx, permuted_outputs_idx], 'out_idxes': [dw_idx], 'inp_shapes': [permuted_inputs_shape, permuted_outputs_shape], 'out_shapes': [dw_shape], 'mask': [], } transcript.append(dw_conv) config.weights_update = dw_idx permutation = [3, 1, 2, 0] permutation_weights_idx = config.new_tensor() permutation_weights_shape = [weights_shape[p] for p in permutation] permute_weights = { 'layer_type': 'Permute', 'params': permutation, 'inp_idxes': [weights_idx], 'out_idxes': [permutation_weights_idx], 'inp_shapes': [weights_shape], 'out_shapes': [permutation_weights_shape], 'mask': [], } transcript.append(permute_weights) rotated_weights_idx, rotated_weights_shape = config.new_tensor(), permutation_weights_shape rotate_layer = { 'layer_type': 'Rotate', 'params': [1, 2], 'inp_idxes': [permutation_weights_idx], 'out_idxes': [rotated_weights_idx], 'inp_shapes': [permutation_weights_shape], 'out_shapes': [rotated_weights_shape], 'mask': [], } transcript.append(rotate_layer) padded_gradients_idx, padded_gradients_shape = config.new_tensor(), output_shape padded_gradients_shape[1] += (rotated_weights_shape[1] - 1) * 2 padded_gradients_shape[2] += (rotated_weights_shape[2] - 1) * 2 pad_layer = { 'layer_type': 'Pad', 'params': [ 0, 0, rotated_weights_shape[1] - 1, rotated_weights_shape[1] - 1, rotated_weights_shape[2] - 1, rotated_weights_shape[2] - 1, 0, 0 ], 'inp_idxes': [config.gradient_tensor_idx(output_idx)], 'out_idxes': [padded_gradients_idx], 'inp_shapes': [], 'out_shapes': [], 'mask': [], } transcript.append(pad_layer) dx_idx, dx_shape = config.new_gradient_tensor(inputs_idx), inputs_shape input_conv_layer = { 'layer_type': 'Conv2D', 'params': [CONV2D, VALID, NO_ACTIVATION, self.stride_h, self.stride_w], 'inp_idxes': [padded_gradients_idx, rotated_weights_idx], 'out_idxes': [dx_idx], 'inp_shapes': [padded_gradients_shape, rotated_weights_shape], 'out_shapes': [dx_shape], 'mask': [], } transcript.append(input_conv_layer) permutation = [3, 1, 2, 0] permuted_dw_idx = config.new_tensor() permuted_dw_shape = [dw_shape[p] for p in permutation] permute_dw = { 'layer_type': 'Permute', 'params': permutation, 'inp_idxes': [dw_idx], 'out_idxes': [permuted_dw_idx], 'inp_shapes': [dw_shape], 'out_shapes': [permuted_dw_shape], 'mask': [], } transcript.append(permute_dw) updated_weights_idx, updated_weights_shape = config.new_tensor(), dw_shape # Call a layer to update the outputs of the convolution update_weights_layer = { 'layer_type': 'Update', 'params': [], 'inp_idxes': [weights_idx, permuted_dw_idx], 'out_idxes': [updated_weights_idx], 'inp_shapes': [weights_shape, permuted_dw_shape], 'out_shapes': [updated_weights_shape], 'mask': [], } # transcript.append(update_weights_layer) class Softmax(): def __init__(self, layer): return # TODO: Make this generalizable to all neural networks # (do not assume that softmax is the last layer, fused with CE-loss) def backward(self, layer, transcript, config): sub_layer = { 'layer_type': 'Sub', 'params': [], # y_hat - y 'inp_idxes': [layer['out_idxes'][0], config.label_tensor_idx], 'out_idxes': [config.new_gradient_tensor(layer['inp_idxes'][0])], 'inp_shapes': [layer['out_shapes'][0], layer['out_shapes'][0]], 'out_shapes': [layer['out_shapes'][0]], 'mask': [], } transcript.append(sub_layer) class AveragePool2D(): def __init__(self, layer): return def backward(self, layer, transcript, config): # TODO: This is very model specific, must rewrite to be accurate # We just broadcast dx across 3 axes # 1 x 3 x 3 x 1 -> 1 x 1 x 1 x 1280 div_idx = config.new_tensor() reshape_layer = { 'layer_type': 'Broadcast', 'params': [], 'inp_idxes': [config.gradient_tensor_idx(layer['out_idxes'][0])], 'out_idxes': [div_idx], 'inp_shapes': [layer['out_shapes'][0]], 'out_shapes': [layer['inp_shapes'][0]], 'mask': [], } transcript.append(reshape_layer) out_idx = config.new_gradient_tensor(layer['inp_idxes'][0]) out_shape = layer['inp_shapes'][0] div = { 'layer_type': 'Div', 'params': [layer['inp_shapes'][0][1] * layer['inp_shapes'][0][2]], 'inp_idxes': [div_idx], 'out_idxes': [out_idx], 'inp_shapes': [out_shape], 'out_shapes': [out_shape], 'mask': [], } transcript.append(div) class Reshape(): def __init__(self, layer): return def backward(self, layer, transcript, config): reshape_layer = { 'layer_type': 'Reshape', 'params': [], 'inp_idxes': [config.gradient_tensor_idx(layer['out_idxes'][0])], 'out_idxes': [config.new_gradient_tensor(layer['inp_idxes'][0])], 'inp_shapes': [layer['out_shapes'][0]], 'out_shapes': [layer['inp_shapes'][0]], 'mask': [], } transcript.append(reshape_layer) def produce_graph(): # Read msgpack file with open("examples/v2_1.0_224_truncated/model.msgpack", "rb") as data_file: byte_data = data_file.read() model = msgpack.unpackb(byte_data) # TODO: I'm unsure whether the circuit output is always the last indexed tensor softmax_output_index = int(np.max( [[out for out in layer['out_idxes']] for layer in model['layers']] + [[inp for inp in layer['inp_idxes']] for layer in model['layers']] )[0]) circuit_config = CircuitConfig(softmax_output_index + 1) circuit_config.new_label_tensor() transcript = [] for layer in reversed(model['layers']): fetched_layer = None match layer['layer_type']: case "Conv2D": fetched_layer = Conv2D(layer) case "AveragePool2D": fetched_layer = AveragePool2D(layer) case "Softmax": fetched_layer = Softmax(layer) case _: fetched_layer = Reshape(layer) print(layer['layer_type']) fetched_layer.backward(layer, transcript, circuit_config) print('----------------') model['layers'] += transcript model['inp_idxes'].append(circuit_config.label_tensor_idx) model['out_idxes'] = [31] packed = msgpack.packb(model, use_bin_type=True) with open("./examples/train_graph/train.msgpack", 'wb') as f: f.write(packed) print(model.keys()) return model model = produce_graph() print(model.keys()) model['tensors'] = "" print(model['inp_idxes'], model['out_idxes']) ================================================ FILE: python/converter.py ================================================ import argparse import ast from typing import Literal, Union import tensorflow as tf import numpy as np import tflite import msgpack def get_shape(interpreter: tf.lite.Interpreter, tensor_idx): if tensor_idx == -1: return [] tensor = interpreter.get_tensor(tensor_idx) return list(tensor.shape) def handle_numpy_or_literal(inp: Union[np.ndarray, Literal[0]]): if isinstance(inp, int): return np.array([inp]) return inp def get_inputs(op: tflite.Operator): idxes = handle_numpy_or_literal(op.InputsAsNumpy()) idxes = idxes.tolist() idxes = list(filter(lambda x: x != -1, idxes)) return idxes class Converter: def __init__( self, model_path, scale_factor, k, num_cols, num_randoms, use_selectors, commit, expose_output ): self.model_path = model_path self.scale_factor = scale_factor self.k = k self.num_cols = num_cols self.num_randoms = num_randoms self.use_selectors = use_selectors self.commit = commit self.expose_output = expose_output self.interpreter = tf.lite.Interpreter( model_path=self.model_path, experimental_preserve_all_tensors=True ) self.interpreter.allocate_tensors() with open(self.model_path, 'rb') as f: buf = f.read() self.model = tflite.Model.GetRootAsModel(buf, 0) self.graph = self.model.Subgraphs(0) def valid_activations(self): return [ tflite.ActivationFunctionType.NONE, tflite.ActivationFunctionType.RELU, tflite.ActivationFunctionType.RELU6, ] def _convert_add(self, op: tflite.Operator, generated_tensors: set): # Get params op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Add options is None') opt = tflite.AddOptions() opt.Init(op_opt.Bytes, op_opt.Pos) params = [opt.FusedActivationFunction()] # Get inputs inputs = get_inputs(op) print(generated_tensors) print('Add inputs: ', inputs) if len(inputs) != 2: raise RuntimeError('Add must have 2 inputs') # If both tensors are generated, do nothing print(inputs[0] in generated_tensors, inputs[1] in generated_tensors) if (inputs[0] in generated_tensors) and (inputs[1] in generated_tensors): return ('Add', params) nb_generated = (inputs[0] in generated_tensors) + (inputs[1] in generated_tensors) if nb_generated != 1: raise RuntimeError('Add must have 1 generated tensor') # Check if there are any negative infinities const_tensor = self.interpreter.get_tensor(inputs[0]) if inputs[0] not in generated_tensors else self.interpreter.get_tensor(inputs[1]) if np.any(const_tensor == -np.inf): # Ensure that the constant tensor is all -inf and 0 if not np.all(np.logical_or(np.isneginf(const_tensor), const_tensor == 0)): raise RuntimeError('Add constant tensor must be -inf and 0 only') mask = (const_tensor == -np.inf).astype(np.int64) params = [len(mask.shape)] + list(mask.shape) params += mask.flatten().tolist() return ('MaskNegInf', params) else: return ('Add', params) def to_dict(self, start_layer, end_layer): interpreter = self.interpreter model = self.model graph = self.graph if graph is None: raise RuntimeError('Graph is None') input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() for inp_detail in input_details: inp = np.zeros(inp_detail['shape'], dtype=inp_detail['dtype']) interpreter.set_tensor(inp_detail['index'], inp) # for i, inp in enumerate(inps): # interpreter.set_tensor(input_details[i]['index'], inp) interpreter.invoke() # Get layers generated_tensor_idxes = set() for inp in input_details: generated_tensor_idxes.add(inp['index']) layers = [] keep_tensors = set() adjusted_tensors = {} for op_idx in range(graph.OperatorsLength()): op = graph.Operators(op_idx) if op is None: raise RuntimeError('Operator is None') model_opcode = model.OperatorCodes(op.OpcodeIndex()) if model_opcode is None: raise RuntimeError('Operator code is None') op_code = model_opcode.BuiltinCode() # Skip generated tensors for output in handle_numpy_or_literal(op.OutputsAsNumpy()): generated_tensor_idxes.add(output) if op_idx < start_layer: continue if op_idx > end_layer: break # Keep the input tensors for input in handle_numpy_or_literal(op.InputsAsNumpy()): keep_tensors.add(input) # AvgPool2D if op_code == tflite.BuiltinOperator.AVERAGE_POOL_2D: layer_type = 'AveragePool2D' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('AvgPool2D options is None') opt = tflite.Pool2DOptions() opt.Init(op_opt.Bytes, op_opt.Pos) params = [opt.FilterHeight(), opt.FilterWidth(), opt.StrideH(), opt.StrideW()] elif op_code == tflite.BuiltinOperator.MAX_POOL_2D: layer_type = 'MaxPool2D' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('MaxPool2D options is None') opt = tflite.Pool2DOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.Padding() == tflite.Padding.SAME: raise NotImplementedError('SAME padding is not supported') if opt.FusedActivationFunction() != tflite.ActivationFunctionType.NONE: raise NotImplementedError('Fused activation is not supported') params = [opt.FilterHeight(), opt.FilterWidth(), opt.StrideH(), opt.StrideW()] # FIXME: hack for Keras... not sure why this isn't being converted properly elif op_code == tflite.BuiltinOperator.CUSTOM: layer_type = 'Conv2D' activation = 0 weights = self.interpreter.get_tensor(op.Inputs(1)) weights = np.transpose(weights, (3, 0, 1, 2)) weights = (weights * self.scale_factor).round().astype(np.int64) adjusted_tensors[op.Inputs(1)] = weights params = [0, 1, activation, 1, 1] # Conv2D elif op_code == tflite.BuiltinOperator.CONV_2D: layer_type = 'Conv2D' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Conv2D options is None') opt = tflite.Conv2DOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.DilationHFactor() != 1 or opt.DilationWFactor() != 1: raise NotImplementedError('Dilation is not supported') if opt.FusedActivationFunction() not in self.valid_activations(): raise NotImplementedError('Unsupported activation function at layer {op_idx}') # 0 is Conv2D params = \ [0] + \ [opt.Padding()] + \ [opt.FusedActivationFunction()] + \ [opt.StrideH(), opt.StrideW()] # DepthwiseConv2D elif op_code == tflite.BuiltinOperator.DEPTHWISE_CONV_2D: layer_type = 'Conv2D' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('DepthwiseConv2D options is None') opt = tflite.DepthwiseConv2DOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.DilationHFactor() != 1 or opt.DilationWFactor() != 1: raise NotImplementedError('Dilation is not supported') if opt.FusedActivationFunction() not in self.valid_activations(): raise NotImplementedError('Unsupported activation function at layer {op_idx}') # 1 is DepthwiseConv2D params = \ [1] + \ [opt.Padding()] + \ [opt.FusedActivationFunction()] + \ [opt.StrideH(), opt.StrideW()] # Fully connected elif op_code == tflite.BuiltinOperator.FULLY_CONNECTED: layer_type = 'FullyConnected' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Fully connected options is None') opt = tflite.FullyConnectedOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.FusedActivationFunction() not in self.valid_activations(): raise NotImplementedError(f'Unsupported activation function at layer {op_idx}') params = [opt.FusedActivationFunction()] elif op_code == tflite.BuiltinOperator.BATCH_MATMUL: layer_type = 'BatchMatMul' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('BatchMatMul options is None') opt = tflite.BatchMatMulOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.AdjX() is True: raise NotImplementedError('AdjX is not supported') params = [int(opt.AdjX()), int(opt.AdjY())] ## Arithmetic # Add elif op_code == tflite.BuiltinOperator.ADD: layer_type, params = self._convert_add(op, generated_tensor_idxes) # Mul elif op_code == tflite.BuiltinOperator.MUL: layer_type = 'Mul' params = [] # Sub elif op_code == tflite.BuiltinOperator.SUB: sub_val = interpreter.get_tensor(op.Inputs(1)) # TODO: this is a bit of a hack if np.any(np.isin(sub_val, 10000)): layer_type = 'MaskNegInf' mask = (sub_val == 10000).astype(np.int64) params = [len(mask.shape)] + list(mask.shape) params += mask.flatten().tolist() else: layer_type = 'Sub' params = [] # Div elif op_code == tflite.BuiltinOperator.DIV: # Implement division as multiplication by the inverse layer_type = 'Mul' div_val = interpreter.get_tensor(op.Inputs(1)) if type(div_val) != np.float32: raise NotImplementedError('Only support one divisor') adjusted_tensors[op.Inputs(1)] = np.array([(self.scale_factor / div_val).round().astype(np.int64)]) params = [] # Pad elif op_code == tflite.BuiltinOperator.PAD: layer_type = 'Pad' tensor_idx = op.Inputs(1) tensor = interpreter.get_tensor(tensor_idx).flatten().astype(np.int64) params = tensor.tolist() # Softmax elif op_code == tflite.BuiltinOperator.SOFTMAX: layer_type = 'Softmax' # TODO: conditionally determine whether or not to subtract the max # It should depend on the input to the softmax if layers[-1]['layer_type'] == 'MaskNegInf': params = layers[-1]['params'] elif layers[-2]['layer_type'] == 'MaskNegInf': params = layers[-2]['params'] params = [params[0] - 1] + params[2:] else: params = [] # Mean elif op_code == tflite.BuiltinOperator.MEAN: layer_type = 'Mean' inp_shape = interpreter.get_tensor(op.Inputs(0)).shape mean_idxes = interpreter.get_tensor(op.Inputs(1)).flatten().astype(np.int64) if len(mean_idxes) + 2 != len(inp_shape): raise NotImplementedError(f'Only mean over all but one axis is supported: {op_idx}') params = mean_idxes.tolist() elif op_code == tflite.BuiltinOperator.SQUARE: layer_type = 'Square' params = [] # Squared difference elif op_code == tflite.BuiltinOperator.SQUARED_DIFFERENCE: layer_type = 'SquaredDifference' params = [] # Pointwise elif op_code == tflite.BuiltinOperator.RSQRT: layer_type = 'Rsqrt' params = [] elif op_code == tflite.BuiltinOperator.LOGISTIC: layer_type = 'Logistic' params = [] elif op_code == tflite.BuiltinOperator.TANH: layer_type = 'Tanh' params = [] elif op_code == tflite.BuiltinOperator.POW: layer_type = 'Pow' power = interpreter.get_tensor(op.Inputs(1)).flatten().astype(np.float32) if power != 3.: raise NotImplementedError(f'Only support power 3') power = power.round().astype(np.int64) if len(power) != 1: raise NotImplementedError(f'Only scalar power is supported: {op_idx}') params = power.tolist() # The following are no-ops in the sense that they don't change the tensor # However, we need to pass along the right tensors # The param says which input to pass along elif op_code == tflite.BuiltinOperator.SHAPE: layer_type = 'Noop' params = [0] elif op_code == tflite.BuiltinOperator.GATHER: layer_type = 'Noop' params = [0] elif op_code == tflite.BuiltinOperator.REDUCE_PROD: # TODO: not sure if this is in general a no-op layer_type = 'Noop' params = [0] elif op_code == tflite.BuiltinOperator.STRIDED_SLICE: # FIXME: this is not in general a no-op layer_type = 'Noop' params = [0] elif op_code == tflite.BuiltinOperator.BROADCAST_ARGS: layer_type = 'Noop' params = [0] elif op_code == tflite.BuiltinOperator.BROADCAST_TO: layer_type = 'Noop' params = [0] ## Shape elif op_code == tflite.BuiltinOperator.RESHAPE: layer_type = 'Reshape' params = [] elif op_code == tflite.BuiltinOperator.TRANSPOSE: layer_type = 'Transpose' params = get_shape(interpreter, op.Inputs(0)) + interpreter.get_tensor(op.Inputs(1)).flatten().astype(np.int64).tolist() elif op_code == tflite.BuiltinOperator.CONCATENATION: # FIXME: This is not in general a no-op layer_type = 'Concatenation' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Concatenation options is None') opt = tflite.ConcatenationOptions() opt.Init(op_opt.Bytes, op_opt.Pos) params = [opt.Axis()] elif op_code == tflite.BuiltinOperator.PACK: layer_type = 'Pack' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Pack options is None') opt = tflite.PackOptions() opt.Init(op_opt.Bytes, op_opt.Pos) params = [opt.Axis()] if params[0] > 1: raise NotImplementedError(f'Only axis=0,1 supported at layer {op_idx}') elif op_code == tflite.BuiltinOperator.SPLIT: layer_type = 'Split' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('Split options is None') opt = tflite.SplitOptions() opt.Init(op_opt.Bytes, op_opt.Pos) axis = interpreter.get_tensor(op.Inputs(0)).flatten().astype(np.int64)[0] num_splits = opt.NumSplits() inp = interpreter.get_tensor(op.Inputs(1)) if inp.shape[axis] % num_splits != 0: raise NotImplementedError(f'Only equal splits supported at layer {op_idx}') params = [int(axis), num_splits] elif op_code == tflite.BuiltinOperator.SLICE: layer_type = 'Slice' begin = interpreter.get_tensor(op.Inputs(1)).flatten().astype(np.int64).tolist() size = interpreter.get_tensor(op.Inputs(2)).flatten().astype(np.int64).tolist() params = begin + size elif op_code == tflite.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: layer_type = 'ResizeNearestNeighbor' op_opt = op.BuiltinOptions() if op_opt is None: raise RuntimeError('ResizeNearestNeighbor options is None') opt = tflite.ResizeNearestNeighborOptions() opt.Init(op_opt.Bytes, op_opt.Pos) if opt.AlignCorners(): raise NotImplementedError(f'Align corners not supported at layer {op_idx}') if not opt.HalfPixelCenters(): raise NotImplementedError(f'Half pixel centers not supported at layer {op_idx}') # Can take the out shape directly from the tensor params = [int(opt.AlignCorners()), int(opt.HalfPixelCenters())] # Not implemented else: op_name = None for attr in dir(tflite.BuiltinOperator): if not attr.startswith('__'): if getattr(tflite.BuiltinOperator, attr) == op_code: op_name = attr raise NotImplementedError('Unsupported operator at layer {}: {}, {}'.format(op_idx, op_code, op_name)) inp_idxes = get_inputs(op) # FIXME: hack for testing rsqrt_overflows = [99, 158, 194, 253, 289, 348] if op_idx in rsqrt_overflows: if op_code == tflite.BuiltinOperator.RSQRT: mask = [0, 1] else: mask = [] else: mask = [] layers.append({ 'layer_type': layer_type, 'inp_idxes': inp_idxes, 'inp_shapes': [get_shape(interpreter, inp_idx) for inp_idx in inp_idxes], 'out_idxes': [op.Outputs(i) for i in range(op.OutputsLength())], 'out_shapes': [get_shape(interpreter, op.Outputs(i)) for i in range(op.OutputsLength())], 'params': params, 'mask': mask, }) print(layers) print() # Get tensors print('keep tensors:', keep_tensors) tensors = [] for tensor_idx in range(graph.TensorsLength()): if tensor_idx not in keep_tensors: continue tensor = graph.Tensors(tensor_idx) if tensor is None: raise NotImplementedError('Tensor is None') if tensor_idx in generated_tensor_idxes: print(f'skipping generated tensor: {format(tensor_idx)}, {tensor.Name()}') continue shape = [] for i in range(tensor.ShapeLength()): shape.append(int(tensor.Shape(i))) if shape == []: shape = [1] tensor_data = interpreter.get_tensor(tensor_idx) if tensor.Type() == tflite.TensorType.FLOAT32: tensor_data = (tensor_data * self.scale_factor).round().astype(np.int64) elif tensor.Type() == tflite.TensorType.INT32: tensor_data = tensor_data.astype(np.int64) elif tensor.Type() == tflite.TensorType.INT64: continue else: raise NotImplementedError('Unsupported tensor type: {}'.format(tensor.Type())) if tensor_idx in adjusted_tensors: tensor_data = adjusted_tensors[tensor_idx] shape = tensor_data.shape tensors.append({ 'idx': tensor_idx, 'shape': shape, 'data': tensor_data.flatten().tolist(), }) # print(tensor_idx, tensor.Type(), tensor.Name(), tensors[-1]['shape']) # print(np.abs(tensor_data).max()) commit_before = [] commit_after = [] if self.commit: input_tensors = [inp['index'] for inp in input_details] weight_tensors = [tensor['idx'] for tensor in tensors if tensor['idx'] not in input_tensors] commit_before = [weight_tensors, input_tensors] output_tensors = [out['index'] for out in output_details] commit_after = [output_tensors] out_idxes = layers[-1]['out_idxes'] if self.expose_output else [] d = { 'global_sf': self.scale_factor, 'k': self.k, 'num_cols': self.num_cols, 'num_random': self.num_randoms, 'inp_idxes': [inp['index'] for inp in input_details], # 'out_idxes': [out['index'] for out in output_details], 'out_idxes': out_idxes, 'layers': layers, 'tensors': tensors, 'use_selectors': self.use_selectors, 'commit_before': commit_before, 'commit_after': commit_after, } print() print(d['layers'][-1]) # d['out_idxes'] = [14] print(d.keys()) print(d['out_idxes']) return d def to_msgpack(self, start_layer, end_layer, use_selectors=True): d = self.to_dict(start_layer, end_layer) model_packed = msgpack.packb(d, use_bin_type=True) d['tensors'] = [] config_packed = msgpack.packb(d, use_bin_type=True) return model_packed, config_packed def main(): parser = argparse.ArgumentParser() parser.add_argument('--model', type=str, required=True) parser.add_argument('--model_output', type=str, required=True) parser.add_argument('--config_output', type=str, required=True) parser.add_argument('--scale_factor', type=int, default=2**16) parser.add_argument('--k', type=int, default=19) parser.add_argument('--eta', type=float, default=0.001) parser.add_argument('--num_cols', type=int, default=6) parser.add_argument('--use_selectors', action=argparse.BooleanOptionalAction, required=False, default=True) parser.add_argument('--commit', action=argparse.BooleanOptionalAction, required=False, default=False) parser.add_argument('--expose_output', action=argparse.BooleanOptionalAction, required=False, default=True) parser.add_argument('--start_layer', type=int, default=0) parser.add_argument('--end_layer', type=int, default=10000) parser.add_argument('--num_randoms', type=int, default=20001) args = parser.parse_args() converter = Converter( args.model, args.scale_factor, args.k, args.num_cols, args.num_randoms, args.use_selectors, args.commit, args.expose_output, ) model_packed, config_packed = converter.to_msgpack( start_layer=args.start_layer, end_layer=args.end_layer, ) if model_packed is None: raise Exception('Failed to convert model') with open(args.model_output, 'wb') as f: f.write(model_packed) with open(args.config_output, 'wb') as f: f.write(config_packed) if __name__ == '__main__': main() ================================================ FILE: python/input_converter.py ================================================ import argparse import ast import numpy as np import msgpack def main(): parser = argparse.ArgumentParser() parser.add_argument('--model_config', type=str, required=True) parser.add_argument('--inputs', type=str, required=True) parser.add_argument('--output', type=str, required=True) args = parser.parse_args() inputs = args.inputs.split(',') with open(args.model_config, 'rb') as f: model_config = msgpack.unpackb(f.read()) input_idxes = model_config['inp_idxes'] scale_factor = model_config['global_sf'] # Get the input shapes from the layers input_shapes = [[0] for _ in input_idxes] for layer in model_config['layers']: for layer_inp_idx, layer_shape in zip(layer['inp_idxes'], layer['inp_shapes']): for index, inp_idx in enumerate(input_idxes): if layer_inp_idx == inp_idx: input_shapes[index] = layer_shape tensors = [] for inp, shape, idx in zip(inputs, input_shapes, input_idxes): tensor = np.load(inp).reshape(shape) tensor = (tensor * scale_factor).round().astype(np.int64) tensors.append({ 'idx': idx, 'shape': shape, 'data': tensor.flatten().tolist(), }) packed = msgpack.packb(tensors, use_bin_type=True) with open(args.output, 'wb') as f: f.write(packed) if __name__ == '__main__': main() ================================================ FILE: python/training_converter.py ================================================ # A converter for training data # Performs the conversion npy -> msgpack # TODO: Ensure that training works with models that take in multiple input shapes # # Shortcut: # `python3 python/training_converter.py --input_shapes 7,7,320 --input_idxes 1,0 --output training_data/inputs.msgpack --labels_output training_data/labels.msgpack` # import argparse import ast import numpy as np import msgpack import os NUM_LOADS = 1 SF = 1 << 17 def main(): parser = argparse.ArgumentParser() parser.add_argument('--input_shapes', type=str, required=True) parser.add_argument('--output', type=str, required=True) TRAINING_DIRECTORY = './testing/data/pre_last_conv/flowers/train' args = parser.parse_args() input_shapes = ast.literal_eval(args.input_shapes) loaded = 0 tensors = [] num_classes = os.listdir(TRAINING_DIRECTORY) first_file = "0.npy" for file_name in os.listdir(TRAINING_DIRECTORY): if loaded == NUM_LOADS: break label = int(first_file[:-4]) data_array = np.load(TRAINING_DIRECTORY + '/' + first_file) input_shape = input_shapes for idx in range(data_array.shape[0]): print(SF) print((np.vstack(data_array) * SF).round().astype(np.int64)) tensors.append({ 'idx': 0, 'shape': input_shape, 'data': list(map(lambda x: int(x), list((data_array[idx] * SF).round().astype(np.int64).flatten()))), }) # represent the label as a one hot encoding one_hot = np.zeros(102) one_hot[label] = SF print("IMPORTANT LABEL", label) print("IMPORTANT LABEL", data_array[idx].flatten()[:500]) # print(one_hot.shape()) tensors.append({ 'idx': 11, 'shape': (1, 102), 'data': list(map(lambda x: int(x), one_hot)), }) loaded += 1 if loaded == NUM_LOADS: break packed_inputs = msgpack.packb(tensors, use_bin_type=True) # print(tensors) with open(args.output, 'wb') as f: f.write(packed_inputs) if __name__ == '__main__': main() ================================================ FILE: rustfmt.toml ================================================ tab_spaces = 2 max_width = 100 ================================================ FILE: src/bin/test_circuit.rs ================================================ use halo2_proofs::{dev::MockProver, halo2curves::bn256::Fr}; use zkml::{ model::ModelCircuit, utils::{ helpers::get_public_values, loader::{load_model_msgpack, ModelMsgpack}, }, }; fn main() { let config_fname = std::env::args().nth(1).expect("config file path"); let inp_fname = std::env::args().nth(2).expect("input file path"); let config: ModelMsgpack = load_model_msgpack(&config_fname, &inp_fname); let circuit = ModelCircuit::::generate_from_file(&config_fname, &inp_fname); let _prover = MockProver::run(config.k.try_into().unwrap(), &circuit, vec![vec![]]).unwrap(); let public_vals = get_public_values(); let prover = MockProver::run(config.k.try_into().unwrap(), &circuit, vec![public_vals]).unwrap(); assert_eq!(prover.verify(), Ok(())); } ================================================ FILE: src/bin/time_circuit.rs ================================================ use halo2_proofs::halo2curves::{bn256::Fr, pasta::Fp}; use zkml::{ model::ModelCircuit, utils::{proving_ipa::time_circuit_ipa, proving_kzg::time_circuit_kzg}, }; fn main() { let config_fname = std::env::args().nth(1).expect("config file path"); let inp_fname = std::env::args().nth(2).expect("input file path"); let kzg_or_ipa = std::env::args().nth(3).expect("kzg or ipa"); if kzg_or_ipa != "kzg" && kzg_or_ipa != "ipa" { panic!("Must specify kzg or ipa"); } if kzg_or_ipa == "kzg" { let circuit = ModelCircuit::::generate_from_file(&config_fname, &inp_fname); time_circuit_kzg(circuit); } else { let circuit = ModelCircuit::::generate_from_file(&config_fname, &inp_fname); time_circuit_ipa(circuit); } } ================================================ FILE: src/bin/verify_circuit.rs ================================================ use halo2_proofs::halo2curves::bn256::Fr; use zkml::{ model::ModelCircuit, utils::{loader::load_config_msgpack, proving_kzg::verify_circuit_kzg}, }; fn main() { let config_fname = std::env::args().nth(1).expect("config file path"); let vkey_fname = std::env::args().nth(2).expect("verification key file path"); let proof_fname = std::env::args().nth(3).expect("proof file path"); let public_vals_fname = std::env::args().nth(4).expect("public values file path"); let kzg_or_ipa = std::env::args().nth(5).expect("kzg or ipa"); if kzg_or_ipa != "kzg" && kzg_or_ipa != "ipa" { panic!("Must specify kzg or ipa"); } if kzg_or_ipa == "kzg" { let config = load_config_msgpack(&config_fname); let circuit = ModelCircuit::::generate_from_msgpack(config, false); println!("Loaded configuration"); verify_circuit_kzg(circuit, &vkey_fname, &proof_fname, &public_vals_fname); } else { // Serialization of the verification key doesn't seem to be supported for IPA panic!("Not implemented"); } } ================================================ FILE: src/bin/verify_wav.rs ================================================ use std::fs::File; use halo2_proofs::{dev::MockProver, halo2curves::bn256::Fr}; use zkml::{ model::ModelCircuit, utils::{ helpers::get_public_values, loader::{load_config_msgpack, ModelMsgpack, TensorMsgpack}, }, }; fn main() { let config_fname = std::env::args().nth(1).expect("config file path"); let wav_fname = std::env::args().nth(2).expect("wav file path"); let mut wav_file = File::open(wav_fname).unwrap(); let (_header, data) = wav::read(&mut wav_file).unwrap(); let data = match data { wav::BitDepth::Sixteen(data) => data, _ => panic!("Unsupported bit depth"), }; let data: Vec = data.iter().map(|x| *x as i64).collect(); let base_config = load_config_msgpack(&config_fname); let config = ModelMsgpack { tensors: vec![TensorMsgpack { idx: 0, shape: vec![1, data.len().try_into().unwrap()], data: data, }], inp_idxes: vec![0], out_idxes: vec![], layers: vec![], commit_before: Some(vec![]), commit_after: Some(vec![vec![0]]), ..base_config }; println!("Config: {:?}", config); let k = config.k; let circuit = ModelCircuit::::generate_from_msgpack(config, false); let _prover = MockProver::run(k.try_into().unwrap(), &circuit, vec![vec![]]).unwrap(); let public_vals: Vec = get_public_values(); println!("Public values: {:?}", public_vals); } ================================================ FILE: src/commitments/commit.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use crate::{gadgets::gadget::GadgetConfig, layers::layer::CellRc}; pub trait Commit { fn commit( &self, layouter: impl Layouter, gadget_config: Rc, constants: &HashMap>, values: &Vec>, blinding: CellRc, ) -> Result>, Error>; } ================================================ FILE: src/commitments/packer.rs ================================================ use std::{ cmp::{max, min}, collections::{BTreeMap, HashMap}, marker::PhantomData, rc::Rc, }; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Value}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::layer::{AssignedTensor, CellRc}, }; const NUM_BITS_PER_FIELD_ELEM: usize = 254; pub struct PackerConfig { pub num_bits_per_elem: usize, pub num_elem_per_packed: usize, pub num_packed_per_row: usize, pub exponents: Vec, _marker: PhantomData, } pub struct PackerChip { pub config: PackerConfig, } impl PackerChip { pub fn get_exponents(num_bits_per_elem: usize, num_exponents: usize) -> Vec { let mul_val = F::from(1 << num_bits_per_elem); let mut exponents = vec![F::ONE]; for _ in 1..num_exponents { exponents.push(exponents[exponents.len() - 1] * mul_val); } exponents } pub fn construct(num_bits_per_elem: usize, gadget_config: &GadgetConfig) -> PackerConfig { let columns = &gadget_config.columns; let num_elem_per_packed = if NUM_BITS_PER_FIELD_ELEM / num_bits_per_elem > columns.len() - 1 { columns.len() - 1 } else { // TODO: for many columns, pack many in a single row NUM_BITS_PER_FIELD_ELEM / num_bits_per_elem }; println!("column len: {}", columns.len()); println!("num_bits_per_elem: {}", num_bits_per_elem); println!("NUM_BITS_PER_FIELD_ELEM: {}", NUM_BITS_PER_FIELD_ELEM); println!("num_elem_per_packed: {}", num_elem_per_packed); let num_packed_per_row = max( 1, columns.len() / (num_elem_per_packed * (num_bits_per_elem + 1)), ); println!("num_packed_per_row: {}", num_packed_per_row); let exponents = Self::get_exponents(num_bits_per_elem, num_elem_per_packed); let config = PackerConfig { num_bits_per_elem, num_elem_per_packed, num_packed_per_row, exponents, _marker: PhantomData, }; config } pub fn configure( meta: &mut ConstraintSystem, packer_config: PackerConfig, gadget_config: GadgetConfig, ) -> GadgetConfig { let selector = meta.complex_selector(); let columns = gadget_config.columns; let lookup = gadget_config.tables.get(&GadgetType::InputLookup).unwrap()[0]; let exponents = &packer_config.exponents; let num_bits_per_elem = packer_config.num_bits_per_elem; let shift_val = 1 << (num_bits_per_elem - 1); let shift_val = Expression::Constant(F::from(shift_val as u64)); meta.create_gate("packer", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..packer_config.num_packed_per_row { let offset = i * (packer_config.num_elem_per_packed + 1); let inps = columns[offset..offset + packer_config.num_elem_per_packed] .iter() .map(|col| meta.query_advice(*col, Rotation::cur())) .collect::>(); let outp = meta.query_advice( columns[offset + packer_config.num_elem_per_packed], Rotation::cur(), ); let res = inps .into_iter() .zip(exponents.iter()) .map(|(inp, exp)| (inp + shift_val.clone()) * (*exp)) .fold(Expression::Constant(F::ZERO), |acc, prod| acc + prod); constraints.push(s.clone() * (res - outp)); // constraints.push(s.clone() * Expression::Constant(F::zero())); } constraints }); // Ensure that the weights/inputs are in the correct range for i in 0..packer_config.num_packed_per_row { let offset = i * (packer_config.num_elem_per_packed + 1); for j in 0..packer_config.num_elem_per_packed { meta.lookup("packer lookup", |meta| { let s = meta.query_selector(selector); let inp = meta.query_advice(columns[offset + j], Rotation::cur()); vec![(s * (inp + shift_val.clone()), lookup)] }); } } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::Packer, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } pub fn copy_and_pack_row( &self, mut layouter: impl Layouter, gadget_config: Rc, cells: Vec>, zero: &AssignedCell, ) -> Result>, Error> { let columns = &gadget_config.columns; let selector = gadget_config.selectors.get(&GadgetType::Packer).unwrap()[0]; let num_bits_per_elem = gadget_config.num_bits_per_elem; let shift_val = 1 << (num_bits_per_elem - 1); let shift_val = F::from(shift_val as u64); let outp = layouter.assign_region( || "pack row", |mut region| { if gadget_config.use_selectors { selector.enable(&mut region, 0)?; } let mut packed = vec![]; for i in 0..self.config.num_packed_per_row { let val_offset = i * self.config.num_elem_per_packed; let col_offset = i * (self.config.num_elem_per_packed + 1); let mut vals = cells [val_offset..min(val_offset + self.config.num_elem_per_packed, cells.len())] .iter() .enumerate() .map(|(i, x)| { x.copy_advice(|| "", &mut region, columns[col_offset + i], 0) .unwrap(); x.value().copied() }) .collect::>(); let zero_copied = (cells.len()..self.config.num_elem_per_packed) .map(|i| { zero .copy_advice(|| "", &mut region, columns[col_offset + i], 0) .unwrap(); zero.value().copied() }) .collect::>(); vals.extend(zero_copied); let res = vals.iter().zip(self.config.exponents.iter()).fold( Value::known(F::ZERO), |acc, (inp, exp)| { let res = acc + (*inp + Value::known(shift_val)) * Value::known(*exp); res }, ); let outp = region.assign_advice( || "", columns[col_offset + self.config.num_elem_per_packed], 0, || res, )?; packed.push(Rc::new(outp)); } Ok(packed) }, )?; Ok(outp) } pub fn assign_and_pack_row( &self, mut layouter: impl Layouter, gadget_config: Rc, values: Vec<&F>, zero: &AssignedCell, ) -> Result<(Vec>, Vec>), Error> { let columns = &gadget_config.columns; let selector = gadget_config.selectors.get(&GadgetType::Packer).unwrap()[0]; let num_bits_per_elem = gadget_config.num_bits_per_elem; let shift_val = 1 << (num_bits_per_elem - 1); let shift_val = F::from(shift_val as u64); let outp = layouter.assign_region( || "pack row", |mut region| { if gadget_config.use_selectors { selector.enable(&mut region, 0)?; } let mut packed = vec![]; let mut assigned = vec![]; for i in 0..self.config.num_packed_per_row { let val_offset = i * self.config.num_elem_per_packed; let col_offset = i * (self.config.num_elem_per_packed + 1); let mut values = values [val_offset..min(val_offset + self.config.num_elem_per_packed, values.len())] .iter() .map(|x| **x) .collect::>(); let vals = values .iter() .enumerate() .map(|(i, x)| { let tmp = region .assign_advice(|| "", columns[col_offset + i], 0, || Value::known(*x)) .unwrap(); Rc::new(tmp) }) .collect::>(); assigned.extend(vals); let zero_vals = (values.len()..self.config.num_elem_per_packed) .map(|i| { zero .copy_advice(|| "", &mut region, columns[col_offset + i], 0) .unwrap(); F::ZERO }) .collect::>(); values.extend(zero_vals); let res = values .iter() .zip(self.config.exponents.iter()) .fold(F::ZERO, |acc, (inp, exp)| { let res = acc + (*inp + shift_val) * (*exp); res }); let outp = region.assign_advice( || "", columns[col_offset + self.config.num_elem_per_packed], 0, || Value::known(res), )?; packed.push(Rc::new(outp)); } Ok((packed, assigned)) }, )?; Ok(outp) } pub fn assign_and_pack( &self, mut layouter: impl Layouter, gadget_config: Rc, constants: &HashMap>, tensors: &BTreeMap>, ) -> Result<(BTreeMap>, Vec>), Error> { let mut values = vec![]; for (_, tensor) in tensors { for value in tensor.iter() { values.push(value); } } let mut packed = vec![]; let mut assigned = vec![]; let zero = constants.get(&0).unwrap().clone(); let num_elems_per_row = self.config.num_packed_per_row * self.config.num_elem_per_packed; for i in 0..(values.len().div_ceil(num_elems_per_row)) { let row = values[i * num_elems_per_row..min((i + 1) * num_elems_per_row, values.len())].to_vec(); let (row_packed, row_assigned) = self .assign_and_pack_row( layouter.namespace(|| "pack row"), gadget_config.clone(), row, zero.as_ref(), ) .unwrap(); packed.extend(row_packed); assigned.extend(row_assigned); } let mut assigned_tensors = BTreeMap::new(); let mut start_idx = 0; for (tensor_id, tensor) in tensors { let num_el = tensor.len(); let v = assigned[start_idx..start_idx + num_el].to_vec(); let new_tensor = Array::from_shape_vec(tensor.raw_dim(), v).unwrap(); assigned_tensors.insert(*tensor_id, new_tensor); start_idx += num_el; } Ok((assigned_tensors, packed)) } pub fn copy_and_pack( &self, mut layouter: impl Layouter, gadget_config: Rc, constants: &HashMap>, tensors: &BTreeMap>, ) -> Result>, Error> { let mut values = vec![]; for (_, tensor) in tensors { for value in tensor.iter() { values.push(value.clone()); } } let mut packed = vec![]; let zero = constants.get(&0).unwrap().clone(); let num_elems_per_row = self.config.num_packed_per_row * self.config.num_elem_per_packed; for i in 0..(values.len().div_ceil(num_elems_per_row)) { let row = values[i * num_elems_per_row..min((i + 1) * num_elems_per_row, values.len())].to_vec(); let row_packed = self .copy_and_pack_row( layouter.namespace(|| "pack row"), gadget_config.clone(), row, zero.as_ref(), ) .unwrap(); packed.extend(row_packed); } Ok(packed) } } ================================================ FILE: src/commitments/poseidon_commit.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_gadgets::poseidon::{ primitives::{generate_constants, Absorbing, ConstantLength, Domain, Mds, Spec}, PaddedWord, PoseidonSpongeInstructions, Pow5Chip, Pow5Config, Sponge, }; use halo2_proofs::{ circuit::Layouter, halo2curves::ff::{FromUniformBytes, PrimeField}, plonk::{Advice, Column, ConstraintSystem, Error}, }; use crate::{gadgets::gadget::GadgetConfig, layers::layer::CellRc}; use super::commit::Commit; pub const WIDTH: usize = 3; pub const RATE: usize = 2; pub const L: usize = 8 - WIDTH - 1; #[derive(Clone, Debug)] pub struct PoseidonCommitChip< F: PrimeField + Ord + FromUniformBytes<64>, const WIDTH: usize, const RATE: usize, const L: usize, > { pub poseidon_config: Pow5Config, } #[derive(Debug)] pub struct P128Pow5T3Gen(PhantomData); impl P128Pow5T3Gen { pub fn new() -> Self { P128Pow5T3Gen(PhantomData::default()) } } impl + Ord, const SECURE_MDS: usize> Spec for P128Pow5T3Gen { fn full_rounds() -> usize { 8 } fn partial_rounds() -> usize { 56 } fn sbox(val: F) -> F { val.pow_vartime([5]) } fn secure_mds() -> usize { SECURE_MDS } fn constants() -> (Vec<[F; 3]>, Mds, Mds) { generate_constants::<_, Self, 3, 2>() } } /// A Poseidon hash function, built around a sponge. #[derive(Debug)] pub struct MyHash< F: PrimeField, PoseidonChip: PoseidonSpongeInstructions, S: Spec, D: Domain, const T: usize, const RATE: usize, > { pub sponge: Sponge, RATE>, D, T, RATE>, } impl> PoseidonCommitChip { pub fn configure( meta: &mut ConstraintSystem, // TODO: ?? _input: [Column; L], state: [Column; WIDTH], partial_sbox: Column, ) -> PoseidonCommitChip { let rc_a = (0..WIDTH).map(|_| meta.fixed_column()).collect::>(); let rc_b = (0..WIDTH).map(|_| meta.fixed_column()).collect::>(); meta.enable_constant(rc_b[0]); PoseidonCommitChip { poseidon_config: Pow5Chip::configure::>( meta, state.try_into().unwrap(), partial_sbox, rc_a.try_into().unwrap(), rc_b.try_into().unwrap(), ), } } } impl> Commit for PoseidonCommitChip { fn commit( &self, mut layouter: impl Layouter, _gadget_config: Rc, _constants: &HashMap>, values: &Vec>, blinding: CellRc, ) -> Result>, Error> { let chip = Pow5Chip::construct(self.poseidon_config.clone()); let mut hasher: MyHash, P128Pow5T3Gen, ConstantLength, 3, 2> = Sponge::new(chip, layouter.namespace(|| "sponge")) .map(|sponge| MyHash { sponge }) .unwrap(); let mut new_vals = values .iter() .map(|x| x.clone()) .chain(vec![blinding.clone()]) .collect::>(); while new_vals.len() % L != 0 { new_vals.push(blinding.clone()); } for (i, value) in new_vals .iter() .map(|x| PaddedWord::Message((**x).clone())) .chain( as Domain>::padding(L).map(PaddedWord::Padding)) .enumerate() { hasher .sponge .absorb(layouter.namespace(|| format!("absorb {}", i)), value) .unwrap(); } let outp = hasher .sponge .finish_absorbing(layouter.namespace(|| "finish absorbing")) .unwrap() .squeeze(layouter.namespace(|| "squeeze")) .unwrap(); let outp = Rc::new(outp); Ok(vec![outp]) } } ================================================ FILE: src/commitments.rs ================================================ pub mod commit; pub mod packer; pub mod poseidon_commit; ================================================ FILE: src/gadgets/add_pairs.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type AddPairsConfig = GadgetConfig; pub struct AddPairsChip { config: Rc, _marker: PhantomData, } impl AddPairsChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("add pair", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..columns.len() / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); let res = inp1 + inp2; constraints.append(&mut vec![s.clone() * (res - outp)]) } constraints }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::AddPairs, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } impl Gadget for AddPairsChip { fn name(&self) -> String { "add pairs chip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let inp1 = &vec_inputs[0]; let inp2 = &vec_inputs[1]; assert_eq!(inp1.len(), inp2.len()); let columns = &self.config.columns; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::AddPairs).unwrap()[0]; selector.enable(region, row_offset)?; } let mut outps = vec![]; for i in 0..inp1.len() { let offset = i * self.num_cols_per_op(); let inp1 = inp1[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let inp2 = inp2[i].copy_advice(|| "", region, columns[offset + 1], row_offset)?; let outp = inp1.value().map(|x: &F| x.to_owned()) + inp2.value().map(|x: &F| x.to_owned()); let outp = region.assign_advice(|| "", columns[offset + 2], row_offset, || outp)?; outps.push(outp); } Ok(outps) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp1 = vec_inputs[0].clone(); let mut inp2 = vec_inputs[1].clone(); let initial_len = inp1.len(); while inp1.len() % self.num_inputs_per_row() != 0 { inp1.push(zero); inp2.push(zero); } let vec_inputs = vec![inp1, inp2]; let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/adder.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region, Value}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type AdderConfig = GadgetConfig; pub struct AdderChip { config: Rc, _marker: PhantomData, } impl AdderChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("adder gate", |meta| { let s = meta.query_selector(selector); let gate_inp = columns[0..columns.len() - 1] .iter() .map(|col| meta.query_advice(*col, Rotation::cur())) .collect::>(); let gate_output = meta.query_advice(*columns.last().unwrap(), Rotation::cur()); let res = gate_inp .iter() .fold(Expression::Constant(F::ZERO), |a, b| a + b.clone()); vec![s * (res - gate_output)] }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::Adder, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } // NOTE: The forward pass of the adder adds _everything_ into one cell impl Gadget for AdderChip { fn name(&self) -> String { "adder".to_string() } fn num_cols_per_op(&self) -> usize { self.config.columns.len() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() - 1 } fn num_outputs_per_row(&self) -> usize { 1 } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(vec_inputs.len(), 1); let inp = &vec_inputs[0]; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::Adder).unwrap()[0]; selector.enable(region, row_offset)?; } inp .iter() .enumerate() .map(|(i, cell)| cell.copy_advice(|| "", region, self.config.columns[i], row_offset)) .collect::, _>>()?; let e = inp.iter().fold(Value::known(F::ZERO), |a, b| { a + b.value().map(|x: &F| x.to_owned()) }); let res = region.assign_advice( || "", *self.config.columns.last().unwrap(), row_offset, || e, )?; Ok(vec![res]) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(single_inputs.len(), 1); let mut inputs = vec_inputs[0].clone(); let zero = single_inputs[0].clone(); while inputs.len() % self.num_inputs_per_row() != 0 { inputs.push(&zero); } let mut outputs = self.op_aligned_rows( layouter.namespace(|| "adder forward"), &vec![inputs], single_inputs, )?; while outputs.len() != 1 { while outputs.len() % self.num_inputs_per_row() != 0 { outputs.push(zero.clone()); } let tmp = outputs.iter().map(|x| x).collect::>(); outputs = self.op_aligned_rows( layouter.namespace(|| "adder forward"), &vec![tmp], single_inputs, )?; } Ok(outputs) } } ================================================ FILE: src/gadgets/bias_div_floor_relu6.rs ================================================ use std::{collections::HashMap, marker::PhantomData}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use crate::gadgets::gadget::convert_to_u64; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type BiasDivFloorRelu6Config = GadgetConfig; const SHIFT_MIN_VAL: i64 = -(1 << 30); pub struct BiasDivFloorRelu6Chip { config: BiasDivFloorRelu6Config, _marker: PhantomData, } impl BiasDivFloorRelu6Chip { pub fn construct(config: BiasDivFloorRelu6Config) -> Self { Self { config, _marker: PhantomData, } } pub fn get_map(scale_factor: u64, num_rows: i64, div_outp_min_val: i64) -> HashMap { let div_val = scale_factor; let div_outp_min_val = div_outp_min_val; let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + div_outp_min_val; let val = shifted.clamp(0, 6 * div_val as i64); map.insert(i as i64, val); } map } pub fn num_cols_per_op() -> usize { 5 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.complex_selector(); let sf = Expression::Constant(F::from(gadget_config.scale_factor)); let columns = gadget_config.columns; let mod_lookup = meta.lookup_table_column(); let relu_lookup = meta.lookup_table_column(); let div_lookup = meta.lookup_table_column(); meta.create_gate("bias_mul", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); let inp = meta.query_advice(columns[offset + 0], Rotation::cur()); let bias = meta.query_advice(columns[offset + 1], Rotation::cur()); let div_res = meta.query_advice(columns[offset + 2], Rotation::cur()); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); constraints.push(s.clone() * (inp - (sf.clone() * (div_res - bias) + mod_res))); } constraints }); for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); meta.lookup("bias_div_relu6 lookup", |meta| { let s = meta.query_selector(selector); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); // Constrains that the modulus \in [0, DIV_VAL) vec![(s.clone() * mod_res.clone(), mod_lookup)] }); meta.lookup("bias_div_relu6 lookup", |meta| { let s = meta.query_selector(selector); let div = meta.query_advice(columns[offset + 2], Rotation::cur()); let outp = meta.query_advice(columns[offset + 4], Rotation::cur()); let div_outp_min_val = Expression::Constant(F::from((-SHIFT_MIN_VAL) as u64)); // Constrains that output \in [0, 6 * SF] vec![ (s.clone() * outp, relu_lookup), (s * (div + div_outp_min_val), div_lookup), ] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::BiasDivFloorRelu6, vec![selector]); let mut tables = gadget_config.tables; tables.insert( GadgetType::BiasDivFloorRelu6, vec![mod_lookup, relu_lookup, div_lookup], ); let mut maps = gadget_config.maps; let relu_map = Self::get_map( gadget_config.scale_factor, gadget_config.num_rows as i64, gadget_config.div_outp_min_val, ); maps.insert(GadgetType::BiasDivFloorRelu6, vec![relu_map]); GadgetConfig { columns, selectors, tables, maps, ..gadget_config } } } impl Gadget for BiasDivFloorRelu6Chip { fn name(&self) -> String { "BiasDivRelu6".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let div_val = self.config.scale_factor as i64; let div_outp_min_val_i64 = -self.config.div_outp_min_val; let div_inp_min_val_pos_i64 = -SHIFT_MIN_VAL; let div_inp_min_val_pos = F::from(div_inp_min_val_pos_i64 as u64); let inp = &vec_inputs[0]; let bias = &vec_inputs[1]; assert_eq!(inp.len(), bias.len()); assert_eq!(inp.len() % self.num_inputs_per_row(), 0); let relu_map = &self .config .maps .get(&GadgetType::BiasDivFloorRelu6) .unwrap()[0]; if self.config.use_selectors { let selector = self .config .selectors .get(&GadgetType::BiasDivFloorRelu6) .unwrap()[0]; selector.enable(region, row_offset)?; } let mut outp_cells = vec![]; for (i, (inp, bias)) in inp.iter().zip(bias.iter()).enumerate() { let offset = i * self.num_cols_per_op(); let inp_f = inp.value().map(|x: &F| x.to_owned()); let bias_f = bias.value().map(|x: &F| { let a = *x + div_inp_min_val_pos; let a = convert_to_u64(&a) as i64 - div_inp_min_val_pos_i64; a }); let div_mod_res = inp_f.map(|x: F| { let x_pos = x + div_inp_min_val_pos; let inp = convert_to_u64(&x_pos); // println!("inp: {:?}, bias: {:?}, x_pos: {:?}", inp, bias, x_pos); let div_res = inp as i64 / div_val - (div_inp_min_val_pos_i64 / div_val); let mod_res = inp as i64 % div_val; // println!("div_res: {:?}, mod_res: {:?}", div_res, mod_res); (div_res, mod_res) }); let div_res = div_mod_res.map(|x: (i64, i64)| x.0) + bias_f; let mod_res = div_mod_res.map(|x: (i64, i64)| x.1); let outp = div_res.map(|x: i64| { let mut x_pos = x - div_outp_min_val_i64; if !relu_map.contains_key(&(x_pos)) { println!("x: {}, x_pos: {}", x, x_pos); x_pos = 0; } let outp_val = relu_map.get(&(x_pos)).unwrap(); // println!("x: {}, x_pos: {}, outp_val: {}", x, x_pos, outp_val); F::from(*outp_val as u64) }); // Assign inp, bias inp.copy_advice(|| "", region, self.config.columns[offset + 0], row_offset)?; bias.copy_advice(|| "", region, self.config.columns[offset + 1], row_offset)?; // Assign div_res, mod_res let div_res_cell = region .assign_advice( || "div_res", self.config.columns[offset + 2], row_offset, || { div_res.map(|x: i64| { F::from((x - div_outp_min_val_i64) as u64) - F::from(-div_outp_min_val_i64 as u64) }) }, ) .unwrap(); let _mod_res_cell = region .assign_advice( || "mod_res", self.config.columns[offset + 3], row_offset, || mod_res.map(|x: i64| F::from(x as u64)), ) .unwrap(); let outp_cell = region .assign_advice( || "outp", self.config.columns[offset + 4], row_offset, || outp.map(|x: F| x.to_owned()), ) .unwrap(); // outp_cells.push((outp_cell, div_res_cell)); outp_cells.push(outp_cell); outp_cells.push(div_res_cell); } Ok(outp_cells) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inps = vec_inputs[0].clone(); let mut biases = vec_inputs[1].clone(); // Needed to pad: bias - bias = 0 let default = biases[0].clone(); while inps.len() % self.num_inputs_per_row() != 0 { inps.push(&default); biases.push(&default); } let res = self.op_aligned_rows( layouter.namespace(|| "bias_div_relu6"), &vec![inps, biases], single_inputs, )?; Ok(res) } } ================================================ FILE: src/gadgets/bias_div_round_relu6.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region, Value}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use crate::gadgets::gadget::convert_to_u64; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type BiasDivRoundRelu6Config = GadgetConfig; const NUM_COLS_PER_OP: usize = 5; pub struct BiasDivRoundRelu6Chip { config: Rc, _marker: PhantomData, } impl BiasDivRoundRelu6Chip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn get_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let div_val = scale_factor; let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let val = shifted.clamp(0, 6 * div_val as i64); map.insert(i as i64, val); } map } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.complex_selector(); let sf = Expression::Constant(F::from(gadget_config.scale_factor)); let two = Expression::Constant(F::from(2)); let columns = gadget_config.columns; let mut tables = gadget_config.tables; let div_lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; let relu_lookup = meta.lookup_table_column(); meta.create_gate("bias_mul", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for op_idx in 0..columns.len() / NUM_COLS_PER_OP { let offset = op_idx * NUM_COLS_PER_OP; let inp = meta.query_advice(columns[offset + 0], Rotation::cur()); let bias = meta.query_advice(columns[offset + 1], Rotation::cur()); let div_res = meta.query_advice(columns[offset + 2], Rotation::cur()); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); // ((div - bias) * 2 + mod) * sf = 2 * inp + sf constraints.push( s.clone() * (two.clone() * inp + sf.clone() - (sf.clone() * two.clone() * (div_res - bias) + mod_res)), ); } constraints }); for op_idx in 0..columns.len() / NUM_COLS_PER_OP { let offset = op_idx * NUM_COLS_PER_OP; meta.lookup("bias_div_relu6 lookup", |meta| { let s = meta.query_selector(selector); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); // Constrains that the modulus \in [0, DIV_VAL) // div_val - mod_res \in [0, max_val) vec![(s.clone() * (two.clone() * sf.clone() - mod_res), div_lookup)] }); meta.lookup("bias_div_relu6 lookup", |meta| { let s = meta.query_selector(selector); let div = meta.query_advice(columns[offset + 2], Rotation::cur()); let outp = meta.query_advice(columns[offset + 4], Rotation::cur()); let div_outp_min_val = gadget_config.div_outp_min_val; let div_outp_min_val = Expression::Constant(F::from((-div_outp_min_val) as u64)); // Constrains that output \in [0, 6 * SF] vec![ (s.clone() * (div + div_outp_min_val), div_lookup), (s.clone() * outp, relu_lookup), ] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::BiasDivRoundRelu6, vec![selector]); tables.insert(GadgetType::BiasDivRoundRelu6, vec![relu_lookup]); let mut maps = gadget_config.maps; let relu_map = Self::get_map( gadget_config.scale_factor, gadget_config.min_val, gadget_config.num_rows as i64, ); maps.insert(GadgetType::BiasDivRoundRelu6, vec![relu_map]); GadgetConfig { columns, selectors, tables, maps, ..gadget_config } } } impl Gadget for BiasDivRoundRelu6Chip { fn name(&self) -> String { "BiasDivRelu6".to_string() } fn num_cols_per_op(&self) -> usize { NUM_COLS_PER_OP } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / NUM_COLS_PER_OP } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() * 2 } fn load_lookups(&self, mut layouter: impl Layouter) -> Result<(), Error> { let map = &self.config.maps[&GadgetType::BiasDivRoundRelu6][0]; let relu_lookup = self.config.tables[&GadgetType::BiasDivRoundRelu6][0]; layouter .assign_table( || "bdr round div/relu lookup", |mut table| { for i in 0..self.config.num_rows { let i = i as i64; let val = map.get(&i).unwrap(); table .assign_cell( || "relu lookup", relu_lookup, i as usize, || Value::known(F::from(*val as u64)), ) .unwrap(); } Ok(()) }, ) .unwrap(); Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let div_val = self.config.scale_factor as i64; let div_outp_min_val_i64 = self.config.div_outp_min_val; let div_inp_min_val_pos_i64 = -self.config.shift_min_val; let div_inp_min_val_pos = F::from(div_inp_min_val_pos_i64 as u64); let inp = &vec_inputs[0]; let bias = &vec_inputs[1]; assert_eq!(inp.len(), bias.len()); assert_eq!(inp.len() % self.num_inputs_per_row(), 0); let relu_map = &self .config .maps .get(&GadgetType::BiasDivRoundRelu6) .unwrap()[0]; if self.config.use_selectors { let selector = self .config .selectors .get(&GadgetType::BiasDivRoundRelu6) .unwrap()[0]; selector.enable(region, row_offset).unwrap(); } let mut outp_cells = vec![]; for (i, (inp, bias)) in inp.iter().zip(bias.iter()).enumerate() { let offset = i * NUM_COLS_PER_OP; let inp_f = inp.value().map(|x: &F| x.to_owned()); let bias_f = bias.value().map(|x: &F| { let a = *x + div_inp_min_val_pos; let a = convert_to_u64(&a) as i64 - div_inp_min_val_pos_i64; a }); let div_mod_res = inp_f.map(|x: F| { let x_pos = x + div_inp_min_val_pos; let inp = convert_to_u64(&x_pos) as i64; let div_inp = 2 * inp + div_val; let div_res = div_inp / (2 * div_val) - div_inp_min_val_pos_i64 / div_val; let mod_res = div_inp % (2 * div_val); (div_res, mod_res) }); let div_res = div_mod_res.map(|x: (i64, i64)| x.0) + bias_f; let mod_res = div_mod_res.map(|x: (i64, i64)| x.1); let outp = div_res.map(|x: i64| { let mut x_pos = x - div_outp_min_val_i64; if !relu_map.contains_key(&(x_pos)) { println!("x: {}, x_pos: {}", x, x_pos); x_pos = 0; } let outp_val = relu_map.get(&(x_pos)).unwrap(); F::from(*outp_val as u64) }); // Assign inp, bias inp .copy_advice(|| "", region, self.config.columns[offset + 0], row_offset) .unwrap(); bias .copy_advice(|| "", region, self.config.columns[offset + 1], row_offset) .unwrap(); // Assign div_res, mod_res let div_res_cell = region .assign_advice( || "div_res", self.config.columns[offset + 2], row_offset, || { div_res.map(|x: i64| { F::from((x - div_outp_min_val_i64) as u64) - F::from(-div_outp_min_val_i64 as u64) }) }, ) .unwrap(); let _mod_res_cell = region .assign_advice( || "mod_res", self.config.columns[offset + 3], row_offset, || mod_res.map(|x: i64| F::from(x as u64)), ) .unwrap(); let outp_cell = region .assign_advice( || "outp", self.config.columns[offset + 4], row_offset, || outp.map(|x: F| x.to_owned()), ) .unwrap(); // outp_cells.push((outp_cell, div_res_cell)); outp_cells.push(outp_cell); outp_cells.push(div_res_cell); } Ok(outp_cells) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inps = vec_inputs[0].clone(); let mut biases = vec_inputs[1].clone(); let initial_len = inps.len(); // Needed to pad: bias - bias = 0 let default = biases[0].clone(); while inps.len() % self.num_inputs_per_row() != 0 { inps.push(&default); biases.push(&default); } let res = self .op_aligned_rows( layouter.namespace(|| "bias_div_relu6"), &vec![inps, biases], single_inputs, ) .unwrap(); Ok(res[0..initial_len * 2].to_vec()) } } ================================================ FILE: src/gadgets/dot_prod.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{Advice, Column, ConstraintSystem, Error, Expression}, poly::Rotation, }; use crate::gadgets::adder::AdderChip; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type DotProductConfig = GadgetConfig; pub struct DotProductChip { config: Rc, _marker: PhantomData, } impl DotProductChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn get_input_columns(config: &GadgetConfig) -> Vec> { let num_inputs = (config.columns.len() - 1) / 2; config.columns[0..num_inputs].to_vec() } pub fn get_weight_columns(config: &GadgetConfig) -> Vec> { let num_inputs = (config.columns.len() - 1) / 2; config.columns[num_inputs..config.columns.len() - 1].to_vec() } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = &gadget_config.columns; meta.create_gate("dot product gate", |meta| { let s = meta.query_selector(selector); let gate_inp = DotProductChip::::get_input_columns(&gadget_config) .iter() .map(|col| meta.query_advice(*col, Rotation::cur())) .collect::>(); let gate_weights = DotProductChip::::get_weight_columns(&gadget_config) .iter() .map(|col| meta.query_advice(*col, Rotation::cur())) .collect::>(); let gate_output = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); let res = gate_inp .iter() .zip(gate_weights) .map(|(a, b)| a.clone() * b.clone()) .fold(Expression::Constant(F::ZERO), |a, b| a + b); vec![s * (res - gate_output)] }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::DotProduct, vec![selector]); GadgetConfig { columns: gadget_config.columns, selectors, ..gadget_config } } } impl Gadget for DotProductChip { fn name(&self) -> String { "dot product".to_string() } fn num_cols_per_op(&self) -> usize { self.config.columns.len() } fn num_inputs_per_row(&self) -> usize { (self.config.columns.len() - 1) / 2 } fn num_outputs_per_row(&self) -> usize { 1 } // The caller is expected to pad the inputs fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(vec_inputs.len(), 2); let inp = &vec_inputs[0]; let weights = &vec_inputs[1]; assert_eq!(inp.len(), weights.len()); assert_eq!(inp.len(), self.num_inputs_per_row()); let zero = &single_inputs[0]; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::DotProduct).unwrap()[0]; selector.enable(region, row_offset).unwrap(); } let inp_cols = DotProductChip::::get_input_columns(&self.config); inp .iter() .enumerate() .map(|(i, cell)| cell.copy_advice(|| "", region, inp_cols[i], row_offset)) .collect::, _>>() .unwrap(); let weight_cols = DotProductChip::::get_weight_columns(&self.config); weights .iter() .enumerate() .map(|(i, cell)| cell.copy_advice(|| "", region, weight_cols[i], row_offset)) .collect::, _>>() .unwrap(); // All columns need to be assigned if self.config.columns.len() % 2 == 0 { zero .copy_advice( || "", region, self.config.columns[self.config.columns.len() - 2], row_offset, ) .unwrap(); } let e = inp .iter() .zip(weights.iter()) .map(|(a, b)| a.value().map(|x: &F| *x) * b.value()) .reduce(|a, b| a + b) .unwrap(); let res = region .assign_advice( || "", self.config.columns[self.config.columns.len() - 1], row_offset, || e, ) .unwrap(); Ok(vec![res]) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(vec_inputs.len(), 2); assert_eq!(single_inputs.len(), 1); let zero = &single_inputs[0]; let mut inputs = vec_inputs[0].clone(); let mut weights = vec_inputs[1].clone(); while inputs.len() % self.num_inputs_per_row() != 0 { inputs.push(&zero); weights.push(&zero); } let outputs = layouter .assign_region( || "dot prod rows", |mut region| { let mut outputs = vec![]; for i in 0..inputs.len() / self.num_inputs_per_row() { let inp = inputs[i * self.num_inputs_per_row()..(i + 1) * self.num_inputs_per_row()].to_vec(); let weights = weights[i * self.num_inputs_per_row()..(i + 1) * self.num_inputs_per_row()].to_vec(); let res = self .op_row_region(&mut region, i, &vec![inp, weights], &vec![zero.clone()]) .unwrap(); outputs.push(res[0].clone()); } Ok(outputs) }, ) .unwrap(); let adder_chip = AdderChip::::construct(self.config.clone()); let tmp = outputs.iter().map(|x| x).collect::>(); Ok( adder_chip .forward( layouter.namespace(|| "dot prod adder"), &vec![tmp], single_inputs, ) .unwrap(), ) } } ================================================ FILE: src/gadgets/gadget.rs ================================================ use std::{ collections::{BTreeSet, HashMap}, sync::Arc, }; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::group::ff::PrimeField, plonk::{Advice, Column, Error, Fixed, Selector, TableColumn}, }; use num_bigint::{BigUint, ToBigUint}; use num_traits::cast::ToPrimitive; #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, PartialOrd, Ord)] pub enum GadgetType { AddPairs, Adder, BiasDivRoundRelu6, BiasDivFloorRelu6, DotProduct, Exp, Logistic, Max, Pow, Relu, Rsqrt, Sqrt, SqrtBig, Square, SquaredDiff, SubPairs, Tanh, MulPairs, VarDivRound, VarDivRoundBig, VarDivRoundBig3, Packer, // This is a special case InputLookup, // Dummy placeholder for the input lookup Update, } #[derive(Clone, Debug, Default)] pub struct GadgetConfig { pub used_gadgets: Arc>, pub columns: Vec>, pub fixed_columns: Vec>, pub selectors: HashMap>, pub tables: HashMap>, pub maps: HashMap>>, pub scale_factor: u64, pub shift_min_val: i64, // MUST be divisible by 2 * scale_factor pub num_rows: usize, pub num_cols: usize, pub k: usize, pub eta: f64, pub min_val: i64, pub max_val: i64, pub div_outp_min_val: i64, pub use_selectors: bool, pub commit_before: Vec>, pub commit_after: Vec>, pub num_bits_per_elem: i64, } // TODO: refactor pub fn convert_to_u64(x: &F) -> u64 { let big = BigUint::from_bytes_le(x.to_repr().as_ref()); let big_digits = big.to_u64_digits(); if big_digits.len() > 2 { println!("big_digits: {:?}", big_digits); } if big_digits.len() == 1 { big_digits[0] as u64 } else if big_digits.len() == 0 { 0 } else { panic!(); } } pub fn convert_to_u128(x: &F) -> u128 { let big = BigUint::from_bytes_le(x.to_repr().as_ref()); big.to_biguint().unwrap().to_u128().unwrap() } pub trait Gadget { fn name(&self) -> String; fn num_cols_per_op(&self) -> usize; fn num_inputs_per_row(&self) -> usize; fn num_outputs_per_row(&self) -> usize; fn load_lookups(&self, _layouter: impl Layouter) -> Result<(), Error> { Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error>; // The caller is required to ensure that the inputs are of the correct length. fn op_aligned_rows( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { // Sanity check inputs for inp in vec_inputs.iter() { assert_eq!(inp.len() % self.num_inputs_per_row(), 0); } let outputs = layouter.assign_region( || format!("gadget {}", self.name()), |mut region| { let mut outputs = vec![]; for i in 0..vec_inputs[0].len() / self.num_inputs_per_row() { let mut vec_inputs_row = vec![]; for inp in vec_inputs.iter() { vec_inputs_row.push( inp[i * self.num_inputs_per_row()..(i + 1) * self.num_inputs_per_row()].to_vec(), ); } let row_outputs = self.op_row_region(&mut region, i, &vec_inputs_row, &single_inputs)?; assert_eq!(row_outputs.len(), self.num_outputs_per_row()); outputs.extend(row_outputs); } Ok(outputs) }, )?; Ok(outputs) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), vec_inputs, single_inputs, ) } } ================================================ FILE: src/gadgets/input_lookup.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region, Value}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; pub struct InputLookupChip { config: Rc, _marker: PhantomData, } impl InputLookupChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let lookup = meta.lookup_table_column(); let mut tables = gadget_config.tables; tables.insert(GadgetType::InputLookup, vec![lookup]); GadgetConfig { tables, ..gadget_config } } } impl Gadget for InputLookupChip { fn load_lookups(&self, mut layouter: impl Layouter) -> Result<(), Error> { let lookup = self.config.tables[&GadgetType::InputLookup][0]; layouter .assign_table( || "input lookup", |mut table| { for i in 0..self.config.num_rows as i64 { table .assign_cell( || "mod lookup", lookup, i as usize, || Value::known(F::from(i as u64)), ) .unwrap(); } Ok(()) }, ) .unwrap(); Ok(()) } fn name(&self) -> String { panic!("InputLookupChip should not be called directly") } fn num_cols_per_op(&self) -> usize { panic!("InputLookupChip should not be called directly") } fn num_inputs_per_row(&self) -> usize { panic!("InputLookupChip should not be called directly") } fn num_outputs_per_row(&self) -> usize { panic!("InputLookupChip should not be called directly") } fn op_row_region( &self, _region: &mut Region, _row_offset: usize, _vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { panic!("InputLookupChip should not be called directly") } } ================================================ FILE: src/gadgets/max.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use crate::gadgets::gadget::convert_to_u64; use super::gadget::{Gadget, GadgetConfig, GadgetType}; pub struct MaxChip { config: Rc, _marker: PhantomData, } impl MaxChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.complex_selector(); let columns = gadget_config.columns; let tables = gadget_config.tables; let inp_lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; meta.create_gate("max arithmetic", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..columns.len() / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); constraints.push(s.clone() * (inp1 - outp.clone()) * (inp2 - outp)) } constraints }); for idx in 0..columns.len() / Self::num_cols_per_op() { meta.lookup("max inp1", |meta| { let s = meta.query_selector(selector); let offset = idx * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s * (outp - inp1), inp_lookup)] }); meta.lookup("max inp2", |meta| { let s = meta.query_selector(selector); let offset = idx * Self::num_cols_per_op(); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s * (outp - inp2), inp_lookup)] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::Max, vec![selector]); GadgetConfig { columns, selectors, tables, ..gadget_config } } } impl Gadget for MaxChip { fn name(&self) -> String { "max".to_string() } fn num_cols_per_op(&self) -> usize { 3 } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() * 2 } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(vec_inputs.len(), 1); let inp = &vec_inputs[0]; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::Max).unwrap()[0]; selector.enable(region, row_offset)?; } let min_val_pos = F::from((-self.config.shift_min_val) as u64); let mut outp = vec![]; let chunks: Vec<&[&AssignedCell]> = inp.chunks(self.num_outputs_per_row()).collect(); let i1 = chunks[0]; let i2 = chunks[1]; for (idx, (inp1, inp2)) in i1.iter().zip(i2.iter()).enumerate() { let offset = idx * self.num_cols_per_op(); inp1 .copy_advice(|| "", region, self.config.columns[offset + 0], row_offset) .unwrap(); inp2 .copy_advice(|| "", region, self.config.columns[offset + 1], row_offset) .unwrap(); let max = inp1.value().zip(inp2.value()).map(|(a, b)| { let a = convert_to_u64(&(*a + min_val_pos)); let b = convert_to_u64(&(*b + min_val_pos)); let max = a.max(b); let max = F::from(max) - min_val_pos; max }); let res = region .assign_advice(|| "", self.config.columns[offset + 2], row_offset, || max) .unwrap(); outp.push(res); } Ok(outp) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inputs = vec_inputs[0].clone(); let first = inputs[0]; while inputs.len() % self.num_inputs_per_row() != 0 { inputs.push(first); } // TODO: pretty sure this is correct but check let num_iters = inputs.len().div_ceil(self.num_inputs_per_row()) + self.num_inputs_per_row(); let mut outputs = self.op_aligned_rows( layouter.namespace(|| "max forward"), &vec![inputs], single_inputs, )?; for _ in 0..num_iters { while outputs.len() % self.num_inputs_per_row() != 0 { outputs.push(first.clone()); } let tmp = outputs.iter().map(|x| x).collect::>(); outputs = self.op_aligned_rows( layouter.namespace(|| "max forward"), &vec![tmp], single_inputs, )?; } outputs = vec![outputs.into_iter().next().unwrap()]; Ok(outputs) } } ================================================ FILE: src/gadgets/mul_pairs.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type MulPairsConfig = GadgetConfig; pub struct MulPairsChip { config: Rc, _marker: PhantomData, } impl MulPairsChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("mul pair", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..columns.len() / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); let res = inp1 * inp2; constraints.append(&mut vec![s.clone() * (res - outp)]) } constraints }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::MulPairs, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } impl Gadget for MulPairsChip { fn name(&self) -> String { "MulPairs".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } // TODO: This + below is basically copied from add pairs - make arithmetic generic fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let inp1 = &vec_inputs[0]; let inp2 = &vec_inputs[1]; assert_eq!(inp1.len(), inp2.len()); let columns = &self.config.columns; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::MulPairs).unwrap()[0]; selector.enable(region, row_offset)?; } let mut outps = vec![]; for i in 0..inp1.len() { let offset = i * self.num_cols_per_op(); let inp1 = inp1[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let inp2 = inp2[i].copy_advice(|| "", region, columns[offset + 1], row_offset)?; let outp = inp1.value().map(|x: &F| x.to_owned()) * inp2.value().map(|x: &F| x.to_owned()); let outp = region.assign_advice(|| "", columns[offset + 2], row_offset, || outp)?; outps.push(outp); } Ok(outps) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp1 = vec_inputs[0].clone(); let mut inp2 = vec_inputs[1].clone(); let initial_len = inp1.len(); while inp1.len() % self.num_inputs_per_row() != 0 { inp1.push(zero); inp2.push(zero); } let vec_inputs = vec![inp1, inp2]; let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/nonlinear/exp.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; type ExpGadgetConfig = GadgetConfig; // IMPORTANT: this return exp(x) * SF pub struct ExpGadgetChip { config: Rc, _marker: PhantomData, } impl ExpGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Exp) } } impl NonLinearGadget for ExpGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / (scale_factor as f64); let exp = x.exp(); let exp = (exp * ((scale_factor * scale_factor) as f64)).round() as i64; map.insert(i as i64, exp); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Exp).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Exp).unwrap()[0] } } impl Gadget for ExpGadgetChip { fn name(&self) -> String { "Exp".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Exp)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/logistic.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; pub struct LogisticGadgetChip { config: Rc, _marker: PhantomData, } impl LogisticGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure( meta, gadget_config, GadgetType::Logistic, ) } } impl NonLinearGadget for LogisticGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / (scale_factor as f64); let logistic = 1. / (1. + (-x).exp()); let logistic = (logistic * ((scale_factor) as f64)).round() as i64; map.insert(i as i64, logistic); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Logistic).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Logistic).unwrap()[0] } } impl Gadget for LogisticGadgetChip { fn name(&self) -> String { "LogisticChip".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Logistic)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/non_linearity.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region, Value}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression, Selector}, poly::Rotation, }; use crate::gadgets::gadget::convert_to_u128; use super::super::gadget::Gadget; use super::super::gadget::{GadgetConfig, GadgetType}; const NUM_COLS_PER_OP: usize = 2; pub trait NonLinearGadget: Gadget { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap; fn get_map(&self) -> &HashMap; fn get_selector(&self) -> Selector; fn num_cols_per_op() -> usize { NUM_COLS_PER_OP } fn configure( meta: &mut ConstraintSystem, gadget_config: GadgetConfig, gadget_type: GadgetType, ) -> GadgetConfig { let selector = meta.complex_selector(); let columns = gadget_config.columns; let mut tables = gadget_config.tables; let inp_lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; let outp_lookup = meta.lookup_table_column(); for op_idx in 0..columns.len() / NUM_COLS_PER_OP { let offset = op_idx * NUM_COLS_PER_OP; meta.lookup("non-linear lookup", |meta| { let s = meta.query_selector(selector); let inp = meta.query_advice(columns[offset + 0], Rotation::cur()); let outp = meta.query_advice(columns[offset + 1], Rotation::cur()); let shift_val = gadget_config.min_val; let shift_val_pos = Expression::Constant(F::from((-shift_val) as u64)); vec![ (s.clone() * (inp + shift_val_pos), inp_lookup), (s.clone() * outp, outp_lookup), ] }); } let mut selectors = gadget_config.selectors; selectors.insert(gadget_type, vec![selector]); tables.insert(gadget_type, vec![inp_lookup, outp_lookup]); let mut maps = gadget_config.maps; let non_linear_map = Self::generate_map( gadget_config.scale_factor, gadget_config.min_val, gadget_config.num_rows as i64, ); maps.insert(gadget_type, vec![non_linear_map]); GadgetConfig { columns, selectors, tables, maps, ..gadget_config } } fn load_lookups( &self, mut layouter: impl Layouter, config: Rc, gadget_type: GadgetType, ) -> Result<(), Error> { let map = self.get_map(); let table_col = config.tables.get(&gadget_type).unwrap()[1]; let shift_pos_i64 = -config.shift_min_val; let shift_pos = F::from(shift_pos_i64 as u64); layouter.assign_table( || "non linear table", |mut table| { for i in 0..config.num_rows { let i = i as i64; // FIXME: refactor this let tmp = *map.get(&i).unwrap(); let val = if i == 0 { F::ZERO } else { if tmp >= 0 { F::from(tmp as u64) } else { let tmp = tmp + shift_pos_i64; F::from(tmp as u64) - shift_pos } }; table.assign_cell( || "non linear cell", table_col, i as usize, || Value::known(val), )?; } Ok(()) }, )?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error> { let columns = &gadget_config.columns; let inp = &vec_inputs[0]; let map = self.get_map(); let shift_val_pos_i64 = -gadget_config.shift_min_val; let shift_val_pos = F::from(shift_val_pos_i64 as u64); let min_val = gadget_config.min_val; if gadget_config.use_selectors { let selector = self.get_selector(); selector.enable(region, row_offset)?; } let mut outps = vec![]; for i in 0..inp.len() { let offset = i * 2; inp[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let outp = inp[i].value().map(|x: &F| { let pos = convert_to_u128(&(*x + shift_val_pos)) as i128 - shift_val_pos_i64 as i128; let x = pos as i64 - min_val; let val = *map.get(&x).unwrap(); if x == 0 { F::ZERO } else { if val >= 0 { F::from(val as u64) } else { let val_pos = val + shift_val_pos_i64; F::from(val_pos as u64) - F::from(shift_val_pos_i64 as u64) } } }); let outp = region.assign_advice(|| "nonlinearity", columns[offset + 1], row_offset, || outp)?; outps.push(outp); } Ok(outps) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let inp_len = vec_inputs[0].len(); let mut inp = vec_inputs[0].clone(); while inp.len() % self.num_inputs_per_row() != 0 { inp.push(zero); } let vec_inputs = vec![inp]; let outp = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, &single_inputs, )?; Ok(outp[0..inp_len].to_vec()) } } ================================================ FILE: src/gadgets/nonlinear/pow.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; // IMPORTANT: PowGadget assumes a single power across the entire DAG pub struct PowGadgetChip { config: Rc, _marker: PhantomData, } impl PowGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Pow) } } impl NonLinearGadget for PowGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let power = 3.; // FIXME: need to make this variable somehow... let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / (scale_factor as f64); let y = x.powf(power); let y = (y * ((scale_factor) as f64)).round() as i64; map.insert(i as i64, y); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Pow).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Pow).unwrap()[0] } } impl Gadget for PowGadgetChip { fn name(&self) -> String { "PowGadgetChip".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Pow)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/relu.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; pub struct ReluChip { config: Rc, _marker: PhantomData, } impl ReluChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Relu) } } impl NonLinearGadget for ReluChip { fn generate_map(_scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let relu = shifted.max(0); map.insert(i as i64, relu); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Relu).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Relu).unwrap()[0] } } impl Gadget for ReluChip { fn name(&self) -> String { "Relu".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Relu)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/rsqrt.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; pub struct RsqrtGadgetChip { config: Rc, _marker: PhantomData, } impl RsqrtGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Rsqrt) } } impl NonLinearGadget for RsqrtGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / (scale_factor as f64); let sqrt = x.sqrt(); let rsqrt = 1.0 / sqrt; let rsqrt = (rsqrt * (scale_factor as f64)).round() as i64; map.insert(i as i64, rsqrt); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Rsqrt).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Rsqrt).unwrap()[0] } } impl Gadget for RsqrtGadgetChip { fn name(&self) -> String { "RsqrtGadget".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Rsqrt)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/sqrt.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; pub struct SqrtGadgetChip { config: Rc, _marker: PhantomData, } impl SqrtGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Sqrt) } } impl NonLinearGadget for SqrtGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / (scale_factor as f64); let sqrt = x.sqrt(); let sqrt = (sqrt * (scale_factor as f64)).round() as i64; map.insert(i as i64, sqrt); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Sqrt).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Sqrt).unwrap()[0] } } impl Gadget for SqrtGadgetChip { fn name(&self) -> String { "SqrtGadget".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Sqrt)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear/tanh.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, }; use super::{ super::gadget::{Gadget, GadgetConfig, GadgetType}, non_linearity::NonLinearGadget, }; pub struct TanhGadgetChip { config: Rc, _marker: PhantomData, } impl TanhGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { as NonLinearGadget>::configure(meta, gadget_config, GadgetType::Tanh) } } impl NonLinearGadget for TanhGadgetChip { fn generate_map(scale_factor: u64, min_val: i64, num_rows: i64) -> HashMap { let scale_factor = scale_factor as f64; let mut map = HashMap::new(); for i in 0..num_rows { let shifted = i + min_val; let x = (shifted as f64) / scale_factor; let y = x.tanh(); let y = (y * scale_factor).round() as i64; map.insert(i as i64, y); } map } fn get_map(&self) -> &HashMap { &self.config.maps.get(&GadgetType::Tanh).unwrap()[0] } fn get_selector(&self) -> halo2_proofs::plonk::Selector { self.config.selectors.get(&GadgetType::Tanh).unwrap()[0] } } impl Gadget for TanhGadgetChip { fn name(&self) -> String { "TanhGadgetChip".to_string() } fn num_cols_per_op(&self) -> usize { as NonLinearGadget>::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn load_lookups(&self, layouter: impl Layouter) -> Result<(), Error> { NonLinearGadget::load_lookups(self, layouter, self.config.clone(), GadgetType::Tanh)?; Ok(()) } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::op_row_region( self, region, row_offset, vec_inputs, single_inputs, self.config.clone(), ) } fn forward( &self, layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { NonLinearGadget::forward(self, layouter, vec_inputs, single_inputs) } } ================================================ FILE: src/gadgets/nonlinear.rs ================================================ pub mod exp; pub mod logistic; pub mod non_linearity; pub mod pow; pub mod relu; pub mod rsqrt; pub mod sqrt; pub mod tanh; ================================================ FILE: src/gadgets/sqrt_big.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use crate::gadgets::gadget::convert_to_u64; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type SqrtBigConfig = GadgetConfig; pub struct SqrtBigChip { config: Rc, _marker: PhantomData, } impl SqrtBigChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.complex_selector(); let two = Expression::Constant(F::from(2)); let columns = gadget_config.columns; let tables = gadget_config.tables; let inp_lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; // TODO: prove that these constraints work meta.create_gate("sqrt_big arithm", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); let inp = meta.query_advice(columns[offset + 0], Rotation::cur()); let sqrt = meta.query_advice(columns[offset + 1], Rotation::cur()); let rem = meta.query_advice(columns[offset + 2], Rotation::cur()); let lhs = inp.clone(); let rhs = sqrt.clone() * sqrt.clone() + rem.clone(); constraints.push(s.clone() * (lhs - rhs)); } constraints }); for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); meta.lookup("sqrt_big sqrt lookup", |meta| { let s = meta.query_selector(selector); let sqrt = meta.query_advice(columns[offset + 1], Rotation::cur()); vec![(s.clone() * sqrt, inp_lookup)] }); meta.lookup("sqrt_big rem lookup", |meta| { let s = meta.query_selector(selector); let sqrt = meta.query_advice(columns[offset + 1], Rotation::cur()); let rem = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s.clone() * (rem + sqrt), inp_lookup)] }); meta.lookup("sqrt_big sqrt - rem lookup", |meta| { let s = meta.query_selector(selector); let sqrt = meta.query_advice(columns[offset + 1], Rotation::cur()); let rem = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s.clone() * (two.clone() * sqrt - rem), inp_lookup)] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::SqrtBig, vec![selector]); GadgetConfig { columns, tables, selectors, ..gadget_config } } } impl Gadget for SqrtBigChip { fn name(&self) -> String { "sqrt_big".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let inps = &vec_inputs[0]; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::SqrtBig).unwrap()[0]; selector.enable(region, row_offset)?; } let mut outp_cells = vec![]; for (i, inp) in inps.iter().enumerate() { let offset = i * self.num_cols_per_op(); inp.copy_advice( || "sqrt_big", region, self.config.columns[offset], row_offset, )?; let outp = inp.value().map(|x: &F| { let inp_val = convert_to_u64(x) as i64; let fsqrt = (inp_val as f64).sqrt(); let sqrt = fsqrt.round() as i64; let rem = inp_val - sqrt * sqrt; (sqrt, rem) }); let sqrt_cell = region.assign_advice( || "sqrt_big", self.config.columns[offset + 1], row_offset, || outp.map(|x| F::from(x.0 as u64)), )?; let _rem_cell = region.assign_advice( || "sqrt_big", self.config.columns[offset + 2], row_offset, || { outp.map(|x| { let rem_pos = x.1 + x.0; F::from(rem_pos as u64) - F::from(x.0 as u64) }) }, )?; outp_cells.push(sqrt_cell); } Ok(outp_cells) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp = vec_inputs[0].clone(); let inp_len = inp.len(); while inp.len() % self.num_inputs_per_row() != 0 { inp.push(zero); } let vec_inputs = vec![inp]; let outp = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(outp[0..inp_len].to_vec()) } } ================================================ FILE: src/gadgets/square.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; pub struct SquareGadgetChip { config: Rc, _marker: PhantomData, } impl SquareGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } // TODO: it would be more efficient to do the division here directly pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("square gate", |meta| { let s = meta.query_selector(selector); let gate_inp = meta.query_advice(columns[0], Rotation::cur()); let gate_output = meta.query_advice(columns[1], Rotation::cur()); let res = gate_inp.clone() * gate_inp; vec![s * (res - gate_output)] }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::Square, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } impl Gadget for SquareGadgetChip { fn name(&self) -> String { "SquareChip".to_string() } fn num_cols_per_op(&self) -> usize { 2 } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { assert_eq!(vec_inputs.len(), 1); if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::Square).unwrap()[0]; selector.enable(region, row_offset)?; } let inps = &vec_inputs[0]; let mut outp = vec![]; for (i, inp) in inps.iter().enumerate() { let offset = i * self.num_cols_per_op(); inp.copy_advice(|| "", region, self.config.columns[offset], row_offset)?; let outp_val = inp.value().map(|x: &F| x.to_owned() * x.to_owned()); let outp_cell = region.assign_advice( || "square output", self.config.columns[offset + 1], row_offset, || outp_val, )?; outp.push(outp_cell); } Ok(outp) } fn forward( &self, mut layouter: impl halo2_proofs::circuit::Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp = vec_inputs[0].clone(); let initial_len = inp.len(); while inp.len() % self.num_inputs_per_row() != 0 { inp.push(zero); } let vec_inputs = vec![inp]; let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/squared_diff.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type SquaredDiffConfig = GadgetConfig; pub struct SquaredDiffGadgetChip { config: Rc, _marker: PhantomData, } impl SquaredDiffGadgetChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("squared diff", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..columns.len() / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); let res = (inp1 - inp2).square(); constraints.append(&mut vec![s.clone() * (res - outp)]) } constraints }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::SquaredDiff, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } impl Gadget for SquaredDiffGadgetChip { fn name(&self) -> String { "SquaredDiff".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let inp1 = &vec_inputs[0]; let inp2 = &vec_inputs[1]; assert_eq!(inp1.len(), inp2.len()); let columns = &self.config.columns; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::SquaredDiff).unwrap()[0]; selector.enable(region, row_offset)?; } let mut outps = vec![]; for i in 0..inp1.len() { let offset = i * self.num_cols_per_op(); let inp1 = inp1[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let inp2 = inp2[i].copy_advice(|| "", region, columns[offset + 1], row_offset)?; let outp = inp1.value().map(|x: &F| x.to_owned()) - inp2.value().map(|x: &F| x.to_owned()); let outp = outp * outp; let outp = region.assign_advice(|| "", columns[offset + 2], row_offset, || outp)?; outps.push(outp); } Ok(outps) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp1 = vec_inputs[0].clone(); let mut inp2 = vec_inputs[1].clone(); let initial_len = inp1.len(); while inp1.len() % self.num_inputs_per_row() != 0 { inp1.push(zero); inp2.push(zero); } let vec_inputs = vec![inp1, inp2]; let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/sub_pairs.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error}, poly::Rotation, }; use super::gadget::{Gadget, GadgetConfig, GadgetType}; type SubPairsConfig = GadgetConfig; pub struct SubPairsChip { config: Rc, _marker: PhantomData, } impl SubPairsChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let selector = meta.selector(); let columns = gadget_config.columns; meta.create_gate("sub pair", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; for i in 0..columns.len() / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let inp1 = meta.query_advice(columns[offset + 0], Rotation::cur()); let inp2 = meta.query_advice(columns[offset + 1], Rotation::cur()); let outp = meta.query_advice(columns[offset + 2], Rotation::cur()); let res = inp1 - inp2; constraints.append(&mut vec![s.clone() * (res - outp)]) } constraints }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::SubPairs, vec![selector]); GadgetConfig { columns, selectors, ..gadget_config } } } impl Gadget for SubPairsChip { fn name(&self) -> String { "sub pairs chip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let inp1 = &vec_inputs[0]; let inp2 = &vec_inputs[1]; assert_eq!(inp1.len(), inp2.len()); let columns = &self.config.columns; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::SubPairs).unwrap()[0]; selector.enable(region, row_offset)?; } let mut outps = vec![]; for i in 0..inp1.len() { let offset = i * self.num_cols_per_op(); let inp1 = inp1[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let inp2 = inp2[i].copy_advice(|| "", region, columns[offset + 1], row_offset)?; let outp = inp1.value().map(|x: &F| x.to_owned()) - inp2.value().map(|x: &F| x.to_owned()); let outp = region.assign_advice(|| "", columns[offset + 2], row_offset, || outp)?; outps.push(outp); } Ok(outps) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut inp1 = vec_inputs[0].clone(); let mut inp2 = vec_inputs[1].clone(); let initial_len = inp1.len(); while inp1.len() % self.num_inputs_per_row() != 0 { inp1.push(zero); inp2.push(zero); } let vec_inputs = vec![inp1, inp2]; let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec_inputs, single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/update.rs ================================================ use std::marker::PhantomData; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use crate::gadgets::gadget::{convert_to_u64, GadgetConfig}; use super::gadget::{Gadget, GadgetType}; type UpdateConfig = GadgetConfig; #[derive(Clone, Debug)] pub struct UpdateGadgetChip { config: UpdateConfig, _marker: PhantomData, } impl UpdateGadgetChip { pub fn construct(config: UpdateConfig) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 4 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> UpdateConfig { let tables = &gadget_config.tables; let mod_lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; let columns = gadget_config.columns; let selector = meta.complex_selector(); let div_val = gadget_config.scale_factor; let eta: u64 = (gadget_config.scale_factor as f64 * gadget_config.eta) as u64; meta.create_gate("updater_arith", |meta| { let s = meta.query_selector(selector); let sf = Expression::Constant(F::from(div_val as u64)); let eta = Expression::Constant(F::from(eta as u64)); let mut constraints = vec![]; for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); let w = meta.query_advice(columns[offset], Rotation::cur()); let dw = meta.query_advice(columns[offset + 1], Rotation::cur()); let div = meta.query_advice(columns[offset + 2], Rotation::cur()); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); let expr = (w * sf.clone() - dw * eta.clone()) - (div * sf.clone() + mod_res); constraints.push(s.clone() * expr); } constraints }); for op_idx in 0..columns.len() / Self::num_cols_per_op() { let offset = op_idx * Self::num_cols_per_op(); // Check that mod is smaller than SF meta.lookup("max inp1", |meta| { let s = meta.query_selector(selector); let mod_res = meta.query_advice(columns[offset + 3], Rotation::cur()); // Constrains that the modulus \in [0, DIV_VAL) vec![(s.clone() * mod_res.clone(), mod_lookup)] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::Update, vec![selector]); UpdateConfig { columns, selectors, ..gadget_config } } } impl Gadget for UpdateGadgetChip { fn name(&self) -> String { "updater chip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.config.columns.len() / self.num_cols_per_op() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, _single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let div_val = self.config.scale_factor as i64; let div_val_f = F::from(div_val as u64); let eta = div_val / 1000; let eta = F::from(eta as u64); let div_outp_min_val = self.config.div_outp_min_val; let div_inp_min_val_pos_i64 = -self.config.shift_min_val; let div_inp_min_val_pos = F::from(div_inp_min_val_pos_i64 as u64); let columns = &self.config.columns; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::Update).unwrap()[0]; selector.enable(region, row_offset)?; } let w = &vec_inputs[0]; let dw = &vec_inputs[1]; let mut output_cells = vec![]; for i in 0..w.len() { let offset = i * self.num_cols_per_op(); let _w_cell = w[i].copy_advice(|| "", region, columns[offset + 0], row_offset)?; let _dw_cell = dw[i].copy_advice(|| "", region, columns[offset + 1], row_offset)?; let w_val = w[i].value().map(|x: &F| x.to_owned()); let dw_val = dw[i].value().map(|x: &F| x.to_owned()); let out_scaled = w_val.zip(dw_val).map(|(w, dw)| w * div_val_f - dw * eta); let div_mod = out_scaled.map(|x| { let x_pos = x + div_inp_min_val_pos; let x_pos = if x_pos > F::ZERO { x_pos } else { x_pos + div_val_f }; let inp = convert_to_u64(&x_pos); let div_res = inp as i64 / div_val - (div_inp_min_val_pos_i64 as i64 / div_val); let mod_res = inp as i64 % div_val; (div_res, mod_res) }); let div_res_cell = region .assign_advice( || "div_res", self.config.columns[offset + 2], row_offset, || { div_mod.map(|(x, _): (i64, i64)| { F::from((x - div_outp_min_val as i64) as u64) - F::from(-div_outp_min_val as u64) }) }, ) .unwrap(); let _mod_res_cell = region .assign_advice( || "mod_res", self.config.columns[offset + 3], row_offset, || div_mod.map(|(_, x): (i64, i64)| F::from(x as u64)), ) .unwrap(); output_cells.push(div_res_cell); } Ok(output_cells) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let zero = &single_inputs[0]; let mut w = vec_inputs[0].clone(); let mut dw = vec_inputs[1].clone(); let initial_len = w.len(); while !w.len() % self.num_cols_per_op() == 0 { w.push(zero); } while !dw.len() % self.num_cols_per_op() == 0 { dw.push(zero); } let res = self.op_aligned_rows( layouter.namespace(|| format!("forward row {}", self.name())), &vec![w, dw], single_inputs, )?; Ok(res[0..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/var_div.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use rounded_div::RoundedDiv; use super::gadget::{convert_to_u128, Gadget, GadgetConfig, GadgetType}; type VarDivRoundConfig = GadgetConfig; pub struct VarDivRoundChip { config: Rc, _marker: PhantomData, } impl VarDivRoundChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 3 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let columns = gadget_config.columns; let selector = meta.complex_selector(); let two = Expression::Constant(F::from(2)); let tables = gadget_config.tables; let lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; // a | c | r | ... | b // (2 * a + b) = (2 * b) * c + r // b - r \in [0, 2^N) <-- forces b > r meta.create_gate("var_div_arithm", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; let b = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); let a = meta.query_advice(columns[offset], Rotation::cur()); let c = meta.query_advice(columns[offset + 1], Rotation::cur()); let r = meta.query_advice(columns[offset + 2], Rotation::cur()); let lhs = a.clone() * two.clone() + b.clone(); let rhs = b.clone() * two.clone() * c + r; constraints.push(s.clone() * (lhs - rhs)); } constraints }); for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); // r \in [0, 2^N) meta.lookup("var div range checks r", |meta| { let s = meta.query_selector(selector); let r = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s.clone() * r, lookup)] }); // 2 * b - r \in [0, 2^N) meta.lookup("var div range checks 2b-r", |meta| { let s = meta.query_selector(selector); let b = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); let r = meta.query_advice(columns[offset + 2], Rotation::cur()); vec![(s.clone() * (two.clone() * b - r), lookup)] }); } // b \in [0, 2^N) meta.lookup("var div range checks b", |meta| { let s = meta.query_selector(selector); let b = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); vec![(s.clone() * b, lookup)] }); let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::VarDivRound, vec![selector]); GadgetConfig { columns, tables, selectors, ..gadget_config } } } impl Gadget for VarDivRoundChip { fn name(&self) -> String { "VarDivRoundChip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { (self.config.columns.len() - 1) / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let a_vec = &vec_inputs[0]; // let zero = single_inputs[0].clone(); let b = &single_inputs[1]; let div_outp_min_val_i64 = self.config.div_outp_min_val; let div_inp_min_val_pos_i64 = -self.config.shift_min_val; if self.config.use_selectors { let selector = self.config.selectors.get(&GadgetType::VarDivRound).unwrap()[0]; selector.enable(region, row_offset)?; } b.copy_advice( || "", region, self.config.columns[self.config.columns.len() - 1], row_offset, )?; let mut div_out = vec![]; for (i, a) in a_vec.iter().enumerate() { let offset = i * self.num_cols_per_op(); a.copy_advice(|| "", region, self.config.columns[offset], row_offset)?; let div_mod = a.value().zip(b.value()).map(|(a, b)| { let b = convert_to_u128(b); // Needs to be divisible by b let div_inp_min_val_pos_i64 = div_inp_min_val_pos_i64 / (b as i64) * (b as i64); let div_inp_min_val_pos = F::from(div_inp_min_val_pos_i64 as u64); let a_pos = *a + div_inp_min_val_pos; let a = convert_to_u128(&a_pos); // c = (2 * a + b) / (2 * b) let c_pos = a.rounded_div(b); let c = (c_pos as i128 - (div_inp_min_val_pos_i64 as u128 / b) as i128) as i64; // r = (2 * a + b) % (2 * b) let rem_floor = (a as i128) - (c_pos * b) as i128; let r = 2 * rem_floor + (b as i128); let r = r as i64; (c, r) }); let div_cell = region.assign_advice( || "", self.config.columns[offset + 1], row_offset, || { div_mod.map(|(c, _)| { let offset = F::from(-div_outp_min_val_i64 as u64); let c = F::from((c - div_outp_min_val_i64) as u64); c - offset }) }, )?; let _mod_cell = region.assign_advice( || "", self.config.columns[offset + 2], row_offset, || div_mod.map(|(_, r)| F::from(r as u64)), )?; div_out.push(div_cell); } Ok(div_out) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inps = vec_inputs[0].clone(); let initial_len = inps.len(); // Needed to pad: bias - bias = 0 let default = &single_inputs[0]; while inps.len() % self.num_inputs_per_row() != 0 { inps.push(&default); } let res = self.op_aligned_rows(layouter.namespace(|| "var_div"), &vec![inps], single_inputs)?; Ok(res[..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/var_div_big.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use rounded_div::RoundedDiv; use super::gadget::{convert_to_u128, Gadget, GadgetConfig, GadgetType}; pub struct VarDivRoundBigChip { config: Rc, _marker: PhantomData, } impl VarDivRoundBigChip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 7 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let columns = gadget_config.columns; let selector = meta.complex_selector(); let two = Expression::Constant(F::from(2)); let range = Expression::Constant(F::from(gadget_config.num_rows as u64)); let tables = gadget_config.tables; let lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; // a | c | r | (2 b - r)_1 | (2 b - r)_0 | r_1 | r_0 | ... | b // a / b = c meta.create_gate("var_div_arithm", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; let b = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); // Constrain that (2 * a + b) = (2 * b) * c + r let a = meta.query_advice(columns[offset], Rotation::cur()); let c = meta.query_advice(columns[offset + 1], Rotation::cur()); let r = meta.query_advice(columns[offset + 2], Rotation::cur()); let lhs = a.clone() * two.clone() + b.clone(); let rhs = b.clone() * two.clone() * c + r.clone(); constraints.push(s.clone() * (lhs - rhs)); // Constrain that (2 * b - r) = br1 * max_val + br0 let br1 = meta.query_advice(columns[offset + 3], Rotation::cur()); let br0 = meta.query_advice(columns[offset + 4], Rotation::cur()); let lhs = b.clone() * two.clone() - r.clone(); let rhs = br1 * range.clone() + br0; constraints.push(s.clone() * (lhs - rhs)); // Constrains that r = r1 * max_val + r0 let r1 = meta.query_advice(columns[offset + 5], Rotation::cur()); let r0 = meta.query_advice(columns[offset + 6], Rotation::cur()); let lhs = r.clone(); let rhs = r1 * range.clone() + r0; constraints.push(s.clone() * (lhs - rhs)); } constraints }); // For var div big, we assume that a, b > 0 and are outputs of the previous layer // r must be constrained to be in [0, b) for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); // (2 * b - r)_{1, 0} \in [0, 2^N) meta.lookup("var div big br1", |meta| { let s = meta.query_selector(selector); let br1 = meta.query_advice(columns[offset + 3], Rotation::cur()); vec![(s * br1, lookup)] }); meta.lookup("var div big br0", |meta| { let s = meta.query_selector(selector); let br0 = meta.query_advice(columns[offset + 4], Rotation::cur()); vec![(s * br0, lookup)] }); // r_{1, 0} \in [0, 2^N) meta.lookup("var div big r1", |meta| { let s = meta.query_selector(selector); let r1 = meta.query_advice(columns[offset + 5], Rotation::cur()); vec![(s * r1, lookup)] }); meta.lookup("var div big r0", |meta| { let s = meta.query_selector(selector); let r0 = meta.query_advice(columns[offset + 6], Rotation::cur()); vec![(s * r0, lookup)] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::VarDivRoundBig, vec![selector]); GadgetConfig { columns, tables, selectors, ..gadget_config } } } impl Gadget for VarDivRoundBigChip { fn name(&self) -> String { "VarDivBigRoundChip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { (self.config.columns.len() - 1) / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let a_vec = &vec_inputs[0]; // let zero = single_inputs[0].clone(); let b = &single_inputs[1]; let div_outp_min_val_i64 = self.config.div_outp_min_val; let div_inp_min_val_pos_i64 = -self.config.shift_min_val; let num_rows = self.config.num_rows as i64; if self.config.use_selectors { let selector = self .config .selectors .get(&GadgetType::VarDivRoundBig) .unwrap()[0]; selector.enable(region, row_offset)?; } b.copy_advice( || "", region, self.config.columns[self.config.columns.len() - 1], row_offset, )?; let mut div_out = vec![]; for (i, a) in a_vec.iter().enumerate() { let offset = i * self.num_cols_per_op(); a.copy_advice(|| "", region, self.config.columns[offset], row_offset) .unwrap(); let div_mod = a.value().zip(b.value()).map(|(a, b)| { let b = convert_to_u128(b); // Needs to be divisible by b let div_inp_min_val_pos_i64 = div_inp_min_val_pos_i64 / (b as i64) * (b as i64); let div_inp_min_val_pos = F::from(div_inp_min_val_pos_i64 as u64); let a_pos = *a + div_inp_min_val_pos; let a = convert_to_u128(&a_pos); // c = (2 * a + b) / (2 * b) let c_pos = a.rounded_div(b); let c = (c_pos as i128 - (div_inp_min_val_pos_i64 as u128 / b) as i128) as i64; // r = (2 * a + b) % (2 * b) let rem_floor = (a as i128) - (c_pos * b) as i128; let r = 2 * rem_floor + (b as i128); let r = r as i64; (c, r) }); let br_split = div_mod.zip(b.value()).map(|((_, r), b)| { let b = convert_to_u128(b) as i64; let val = 2 * b - r; let p1 = val / num_rows; let p0 = val % num_rows; // val = p1 * max_val + p0 (p1, p0) }); let r_split = div_mod.map(|(_, r)| { let p1 = r / num_rows; let p0 = r % num_rows; // val = p1 * max_val + p0 (p1, p0) }); let div_cell = region.assign_advice( || "", self.config.columns[offset + 1], row_offset, || { div_mod.map(|(c, _)| { let offset = F::from(-div_outp_min_val_i64 as u64); let c = F::from((c - div_outp_min_val_i64) as u64); c - offset }) }, )?; let _mod_cell = region.assign_advice( || "", self.config.columns[offset + 2], row_offset, || div_mod.map(|(_, r)| F::from(r as u64)), )?; // Assign 2 * b - r to the next 2 columns let _br_split_cell_1 = region.assign_advice( || "", self.config.columns[offset + 3], row_offset, || br_split.map(|(p1, _)| F::from(p1 as u64)), )?; let _br_split_cell_2 = region.assign_advice( || "", self.config.columns[offset + 4], row_offset, || br_split.map(|(_, p0)| F::from(p0 as u64)), )?; // Assign r to the next 2 columns let _r_split_cell_1 = region.assign_advice( || "", self.config.columns[offset + 5], row_offset, || r_split.map(|(p1, _)| F::from(p1 as u64)), )?; let _r_split_cell_2 = region.assign_advice( || "", self.config.columns[offset + 6], row_offset, || r_split.map(|(_, p0)| F::from(p0 as u64)), )?; div_out.push(div_cell); } Ok(div_out) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inps = vec_inputs[0].clone(); let initial_len = inps.len(); // Needed to pad let default = &single_inputs[0]; while inps.len() % self.num_inputs_per_row() != 0 { inps.push(&default); } let res = self.op_aligned_rows( layouter.namespace(|| "var_div_big"), &vec![inps], single_inputs, )?; Ok(res[..initial_len].to_vec()) } } ================================================ FILE: src/gadgets/var_div_big3.rs ================================================ use std::{marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region}, halo2curves::ff::PrimeField, plonk::{ConstraintSystem, Error, Expression}, poly::Rotation, }; use rounded_div::RoundedDiv; use super::gadget::{convert_to_u128, Gadget, GadgetConfig, GadgetType}; pub struct VarDivRoundBig3Chip { config: Rc, _marker: PhantomData, } impl VarDivRoundBig3Chip { pub fn construct(config: Rc) -> Self { Self { config, _marker: PhantomData, } } pub fn num_cols_per_op() -> usize { 9 } pub fn configure(meta: &mut ConstraintSystem, gadget_config: GadgetConfig) -> GadgetConfig { let columns = gadget_config.columns; let selector = meta.complex_selector(); let two = Expression::Constant(F::from(2)); let range = Expression::Constant(F::from(gadget_config.num_rows as u64)); let range_sq = range.clone() * range.clone(); let tables = gadget_config.tables; let lookup = tables.get(&GadgetType::InputLookup).unwrap()[0]; // a | c | r | (2b - r)_2 | (2 b - r)_1 | (2 b - r)_0 | r_2 | r_1 | r_0 | ... | b // a / b = c meta.create_gate("var_div_big3_arithm", |meta| { let s = meta.query_selector(selector); let mut constraints = vec![]; let b = meta.query_advice(columns[columns.len() - 1], Rotation::cur()); for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); // Constrain that (2 * a + b) = (2 * b) * c + r let a = meta.query_advice(columns[offset], Rotation::cur()); let c = meta.query_advice(columns[offset + 1], Rotation::cur()); let r = meta.query_advice(columns[offset + 2], Rotation::cur()); let lhs = a.clone() * two.clone() + b.clone(); let rhs = b.clone() * two.clone() * c + r.clone(); constraints.push(s.clone() * (lhs - rhs)); // Constrain that (2 * b - r) = br1 * max_val + br0 let br2 = meta.query_advice(columns[offset + 3], Rotation::cur()); let br1 = meta.query_advice(columns[offset + 4], Rotation::cur()); let br0 = meta.query_advice(columns[offset + 5], Rotation::cur()); let lhs = b.clone() * two.clone() - r.clone(); let rhs = br2 * range_sq.clone() + br1 * range.clone() + br0; constraints.push(s.clone() * (lhs - rhs)); // Constrains that r = r1 * max_val + r0 let r2 = meta.query_advice(columns[offset + 6], Rotation::cur()); let r1 = meta.query_advice(columns[offset + 7], Rotation::cur()); let r0 = meta.query_advice(columns[offset + 8], Rotation::cur()); let lhs = r.clone(); let rhs = r2 * range_sq.clone() + r1 * range.clone() + r0; constraints.push(s.clone() * (lhs - rhs)); } constraints }); // For var div big, we assume that a, b > 0 and are outputs of the previous layer // r must be constrained to be in [0, b) for i in 0..(columns.len() - 1) / Self::num_cols_per_op() { let offset = i * Self::num_cols_per_op(); // (2 * b - r)_{1, 0} \in [0, 2^N) meta.lookup("var div big br2", |meta| { let s = meta.query_selector(selector); let br2 = meta.query_advice(columns[offset + 3], Rotation::cur()); vec![(s * br2, lookup)] }); meta.lookup("var div big br1", |meta| { let s = meta.query_selector(selector); let br1 = meta.query_advice(columns[offset + 4], Rotation::cur()); vec![(s * br1, lookup)] }); meta.lookup("var div big br0", |meta| { let s = meta.query_selector(selector); let br0 = meta.query_advice(columns[offset + 5], Rotation::cur()); vec![(s * br0, lookup)] }); // r_{1, 0} \in [0, 2^N) meta.lookup("var div big r2", |meta| { let s = meta.query_selector(selector); let r2 = meta.query_advice(columns[offset + 6], Rotation::cur()); vec![(s * r2, lookup)] }); meta.lookup("var div big r1", |meta| { let s = meta.query_selector(selector); let r1 = meta.query_advice(columns[offset + 7], Rotation::cur()); vec![(s * r1, lookup)] }); meta.lookup("var div big r0", |meta| { let s = meta.query_selector(selector); let r0 = meta.query_advice(columns[offset + 8], Rotation::cur()); vec![(s * r0, lookup)] }); } let mut selectors = gadget_config.selectors; selectors.insert(GadgetType::VarDivRoundBig3, vec![selector]); GadgetConfig { columns, tables, selectors, ..gadget_config } } } impl Gadget for VarDivRoundBig3Chip { fn name(&self) -> String { "VarDivBig3RoundChip".to_string() } fn num_cols_per_op(&self) -> usize { Self::num_cols_per_op() } fn num_inputs_per_row(&self) -> usize { (self.config.columns.len() - 1) / self.num_cols_per_op() } fn num_outputs_per_row(&self) -> usize { self.num_inputs_per_row() } fn op_row_region( &self, region: &mut Region, row_offset: usize, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let a_vec = &vec_inputs[0]; // let zero = single_inputs[0].clone(); let b = &single_inputs[1]; let c_shift_base = (-(1_i64 << 62)) as i128; let num_rows = self.config.num_rows as i128; if self.config.use_selectors { let selector = self .config .selectors .get(&GadgetType::VarDivRoundBig3) .unwrap()[0]; selector.enable(region, row_offset)?; } b.copy_advice( || "", region, self.config.columns[self.config.columns.len() - 1], row_offset, )?; let mut div_out = vec![]; for (i, a) in a_vec.iter().enumerate() { let offset = i * self.num_cols_per_op(); a.copy_advice(|| "", region, self.config.columns[offset], row_offset) .unwrap(); let div_mod = a.value().zip(b.value()).map(|(a, b)| { let b = convert_to_u128(b); let c_shift = (-c_shift_base) as u128 / b * b; let div_inp_min_val_pos = F::from(c_shift as u64); let a_pos = *a + div_inp_min_val_pos; let a = convert_to_u128(&a_pos); // c = (2 * a + b) / (2 * b) let c_pos = a.rounded_div(b); let c = c_pos as i128 - (c_shift / b) as i128; // r = (2 * a + b) % (2 * b) let rem_floor = (a as i128) - (c_pos * b) as i128; let r = 2 * rem_floor + (b as i128); (c, r) }); let br_split = div_mod.zip(b.value()).map(|((_, r), b)| { let b = convert_to_u128(b) as i128; let val = 2 * b - r; let p2 = val / (num_rows * num_rows); let p1 = (val % (num_rows * num_rows)) / num_rows; let p0 = val % num_rows; // val = p2 * max_val^2 + p1 * max_val + p0 (p2, p1, p0) }); let r_split = div_mod.map(|(_, r)| { let p2 = r / (num_rows * num_rows); let p1 = (r % (num_rows * num_rows)) / num_rows; let p0 = r % num_rows; // val = p1 * max_val + p0 (p2, p1, p0) }); let div_cell = region.assign_advice( || "", self.config.columns[offset + 1], row_offset, || { div_mod.map(|(c, _)| { let offset = F::from(-c_shift_base as u64); let c = F::from((c - c_shift_base) as u64); c - offset }) }, )?; let _mod_cell = region.assign_advice( || "", self.config.columns[offset + 2], row_offset, || div_mod.map(|(_, r)| F::from(r as u64)), )?; // Assign 2 * b - r to the next 3 columns let _br_split_cell_2 = region.assign_advice( || "", self.config.columns[offset + 3], row_offset, || br_split.map(|(p2, _, _)| F::from(p2 as u64)), )?; let _br_split_cell_1 = region.assign_advice( || "", self.config.columns[offset + 4], row_offset, || br_split.map(|(_, p1, _)| F::from(p1 as u64)), )?; let _br_split_cell_0 = region.assign_advice( || "", self.config.columns[offset + 5], row_offset, || br_split.map(|(_, _, p0)| F::from(p0 as u64)), )?; // Assign r to the next 3 columns let _r_split_cell_2 = region.assign_advice( || "", self.config.columns[offset + 6], row_offset, || r_split.map(|(p2, _, _)| F::from(p2 as u64)), )?; let _r_split_cell_1 = region.assign_advice( || "", self.config.columns[offset + 7], row_offset, || r_split.map(|(_, p1, _)| F::from(p1 as u64)), )?; let _r_split_cell_0 = region.assign_advice( || "", self.config.columns[offset + 8], row_offset, || r_split.map(|(_, _, p0)| F::from(p0 as u64)), )?; div_out.push(div_cell); } Ok(div_out) } fn forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, single_inputs: &Vec<&AssignedCell>, ) -> Result>, Error> { let mut inps = vec_inputs[0].clone(); let initial_len = inps.len(); // Needed to pad let default = &single_inputs[0]; while inps.len() % self.num_inputs_per_row() != 0 { inps.push(&default); } let res = self.op_aligned_rows( layouter.namespace(|| "var_div_big3"), &vec![inps], single_inputs, )?; Ok(res[..initial_len].to_vec()) } } ================================================ FILE: src/gadgets.rs ================================================ pub mod add_pairs; pub mod adder; pub mod bias_div_floor_relu6; pub mod bias_div_round_relu6; pub mod dot_prod; pub mod gadget; pub mod input_lookup; pub mod max; pub mod mul_pairs; pub mod sqrt_big; pub mod square; pub mod squared_diff; pub mod sub_pairs; pub mod update; pub mod var_div; pub mod var_div_big; pub mod var_div_big3; // Generics pub mod nonlinear; ================================================ FILE: src/layers/arithmetic/add.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ add_pairs::AddPairsChip, gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::relu::ReluChip, }, layers::layer::{ActivationType, AssignedTensor, CellRc, GadgetConsumer}, }; use super::{ super::layer::{Layer, LayerConfig}, Arithmetic, }; #[derive(Clone, Debug)] pub struct AddChip {} impl AddChip { fn get_activation(&self, layer_params: &Vec) -> ActivationType { let activation = layer_params[0]; match activation { 0 => ActivationType::None, 1 => ActivationType::Relu, _ => panic!("Unsupported activation type for add"), } } } impl Arithmetic for AddChip { fn gadget_forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, constants: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error> { let add_pairs_chip = AddPairsChip::::construct(gadget_config); let out = add_pairs_chip.forward(layouter.namespace(|| "add chip"), &vec_inputs, constants)?; Ok(out) } } impl Layer for AddChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let activation = self.get_activation(&layer_config.layer_params); // Do the addition let (out, out_shape) = self.arithmetic_forward( layouter.namespace(|| ""), tensors, constants, gadget_config.clone(), )?; // Do the fused activation let out = if activation == ActivationType::Relu { let zero = constants.get(&0).unwrap(); let single_inps = vec![zero.as_ref()]; let out = out.iter().map(|x| x.as_ref()).collect::>(); let relu_chip = ReluChip::::construct(gadget_config); let out = relu_chip.forward(layouter.namespace(|| "relu"), &vec![out], &single_inps)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); out } else if activation == ActivationType::None { out } else { panic!("Unsupported activation type for add"); }; let out = Array::from_shape_vec(IxDyn(out_shape.as_slice()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for AddChip { fn used_gadgets(&self, layer_params: Vec) -> Vec { let activation = self.get_activation(&layer_params); let mut outp = vec![GadgetType::AddPairs]; match activation { ActivationType::Relu => outp.push(GadgetType::Relu), ActivationType::None => (), _ => panic!("Unsupported activation type for add"), } outp } } ================================================ FILE: src/layers/arithmetic/div_var.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, mul_pairs::MulPairsChip, var_div::VarDivRoundChip, }, layers::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer}, }; use super::Arithmetic; pub struct DivVarChip {} // TODO: hack. Used for multiplying by the scale factor impl Arithmetic for DivVarChip { fn gadget_forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, constants: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error> { let mul_pairs_chip = MulPairsChip::::construct(gadget_config.clone()); let out = mul_pairs_chip.forward( layouter.namespace(|| "mul pairs chip"), &vec_inputs, constants, )?; Ok(out) } } impl Layer for DivVarChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &crate::layers::layer::LayerConfig, ) -> Result>, Error> { assert_eq!(tensors.len(), 2); // TODO: We only support dividing by a single number for now assert_eq!(tensors[1].shape().len(), 1); assert_eq!(tensors[1].shape()[0], 1); let sf = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); let sf_tensor = Array::from_shape_vec(IxDyn(&[1]), vec![Rc::new(sf.clone())]).unwrap(); // out = inp * SF let (out, out_shape) = self.arithmetic_forward( layouter.namespace(|| ""), &vec![tensors[0].clone(), sf_tensor], constants, gadget_config.clone(), )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let div = tensors[1].iter().next().unwrap().as_ref(); let zero = constants.get(&0).unwrap().as_ref(); let single_inputs = vec![zero, div]; let out = out.iter().map(|x| x.as_ref()).collect::>(); let out = var_div_chip.forward(layouter.namespace(|| "mul div"), &vec![out], &single_inputs)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(out_shape.as_slice()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for DivVarChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::MulPairs, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/arithmetic/mul.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, mul_pairs::MulPairsChip, var_div::VarDivRoundChip, }, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::{ super::layer::{Layer, LayerConfig}, Arithmetic, }; #[derive(Clone, Debug)] pub struct MulChip {} impl Arithmetic for MulChip { fn gadget_forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, constants: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error> { let mul_pairs_chip = MulPairsChip::::construct(gadget_config.clone()); let out = mul_pairs_chip.forward( layouter.namespace(|| "mul pairs chip"), &vec_inputs, constants, )?; Ok(out) } } // FIXME: move this + add to an arithmetic layer impl Layer for MulChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let (out, out_shape) = self.arithmetic_forward( layouter.namespace(|| ""), tensors, constants, gadget_config.clone(), )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let div = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); let zero = constants.get(&0).unwrap().as_ref(); let single_inputs = vec![zero, div]; let out = out.iter().map(|x| x.as_ref()).collect::>(); let out = var_div_chip.forward(layouter.namespace(|| "mul div"), &vec![out], &single_inputs)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(out_shape.as_slice()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for MulChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::MulPairs, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/arithmetic/sub.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, sub_pairs::SubPairsChip, }, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::{ super::layer::{Layer, LayerConfig}, Arithmetic, }; #[derive(Clone, Debug)] pub struct SubChip {} impl Arithmetic for SubChip { fn gadget_forward( &self, mut layouter: impl Layouter, vec_inputs: &Vec>>, constants: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error> { let sub_pairs_chip = SubPairsChip::::construct(gadget_config); let out = sub_pairs_chip.forward(layouter.namespace(|| "sub chip"), &vec_inputs, constants)?; Ok(out) } } impl Layer for SubChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let (out, out_shape) = self.arithmetic_forward( layouter.namespace(|| ""), tensors, constants, gadget_config.clone(), )?; let out = Array::from_shape_vec(IxDyn(out_shape.as_slice()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for SubChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::SubPairs] } } ================================================ FILE: src/layers/arithmetic.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use crate::{gadgets::gadget::GadgetConfig, utils::helpers::broadcast}; use super::layer::{AssignedTensor, CellRc}; pub mod add; pub mod div_var; pub mod mul; pub mod sub; pub trait Arithmetic { fn gadget_forward( &self, layouter: impl Layouter, vec_inputs: &Vec>>, constants: &Vec<&AssignedCell>, gadget_config: Rc, ) -> Result>, Error>; fn arithmetic_forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, ) -> Result<(Vec>, Vec), Error> { assert_eq!(tensors.len(), 2); // println!("tensors: {:?} {:?}", tensors[0].shape(), tensors[1].shape()); let (inp1, inp2) = broadcast(&tensors[0], &tensors[1]); let out_shape = inp1.shape().clone(); assert_eq!(inp1.shape(), inp2.shape()); let zero = constants.get(&0).unwrap().as_ref(); let inp1_vec = inp1.iter().map(|x| x.as_ref()).collect::>(); let inp2_vec = inp2.iter().map(|x| x.as_ref()).collect::>(); let vec_inputs = vec![inp1_vec, inp2_vec]; let constants = vec![zero]; let out = self.gadget_forward( layouter.namespace(|| ""), &vec_inputs, &constants, gadget_config.clone(), )?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); Ok((out, out_shape.to_vec())) } } ================================================ FILE: src/layers/averager.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use crate::gadgets::gadget::Gadget; use crate::gadgets::{adder::AdderChip, gadget::GadgetConfig, var_div::VarDivRoundChip}; use super::layer::{AssignedTensor, CellRc, LayerConfig}; pub trait Averager { fn splat(&self, input: &AssignedTensor, layer_config: &LayerConfig) -> Vec>>; fn get_div_val( &self, layouter: impl Layouter, tensors: &Vec>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result, Error>; fn avg_forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { // Due to Mean BS // assert_eq!(tensors.len(), 1); let zero = constants.get(&0).unwrap().as_ref(); let inp = &tensors[0]; let splat_inp = self.splat(inp, layer_config); let adder_chip = AdderChip::::construct(gadget_config.clone()); let single_inputs = vec![zero]; let mut added = vec![]; for i in 0..splat_inp.len() { let tmp = splat_inp[i].iter().map(|x| x.as_ref()).collect::>(); let tmp = adder_chip.forward( layouter.namespace(|| format!("average {}", i)), &vec![tmp], &single_inputs, )?; added.push(tmp[0].clone()); } let div = self.get_div_val( layouter.namespace(|| "average div"), tensors, gadget_config.clone(), layer_config, )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let single_inputs = vec![zero, &div]; let added = added.iter().map(|x| x).collect::>(); let dived = var_div_chip.forward( layouter.namespace(|| "average div"), &vec![added], &single_inputs, )?; let dived = dived.into_iter().map(|x| Rc::new(x)).collect::>(); Ok(dived) } } ================================================ FILE: src/layers/avg_pool_2d.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Value}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::max_pool_2d::MaxPool2DChip, }; use super::{ averager::Averager, layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}, }; pub struct AvgPool2DChip {} impl Averager for AvgPool2DChip { fn splat(&self, input: &AssignedTensor, layer_config: &LayerConfig) -> Vec>> { assert_eq!(input.shape().len(), 4); // Don't support batch size > 1 yet assert_eq!(input.shape()[0], 1); // TODO: refactor this MaxPool2DChip::splat(input, layer_config).unwrap() } fn get_div_val( &self, mut layouter: impl Layouter, _tensors: &Vec>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result, Error> { // FIXME: this needs to be revealed let div = layer_config.layer_params[0] * layer_config.layer_params[1]; let div = F::from(div as u64); let div = layouter .assign_region( || "avg pool 2d div", |mut region| { let div = region .assign_advice( || "avg pool 2d div", gadget_config.columns[0], 0, || Value::known(div), ) .unwrap(); Ok(div) }, ) .unwrap(); Ok(div) } } impl Layer for AvgPool2DChip { fn forward( &self, layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let dived = self .avg_forward(layouter, tensors, constants, gadget_config, layer_config) .unwrap(); let inp = &tensors[0]; // TODO: refactor this let out_xy = MaxPool2DChip::shape(inp, layer_config); let out_shape = vec![1, out_xy.0, out_xy.1, inp.shape()[3]]; println!("out_shape: {:?}", out_shape); let out = Array::from_shape_vec(IxDyn(&out_shape), dived).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for AvgPool2DChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::Adder, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/batch_mat_mul.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, Axis, IxDyn}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::fully_connected::FullyConnectedConfig, }; use super::{ fully_connected::FullyConnectedChip, layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}, }; pub struct BatchMatMulChip {} impl Layer for BatchMatMulChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp1 = &tensors[0]; let inp2 = &tensors[1]; println!("inp1: {:?}", inp1.shape()); println!("inp2: {:?}", inp2.shape()); assert_eq!(inp1.ndim(), 3); assert_eq!(inp2.ndim(), 3); assert_eq!(inp1.shape()[0], inp2.shape()[0]); let adj_y = layer_config.layer_params[1] == 1; if adj_y { assert_eq!(inp1.shape()[2], inp2.shape()[2]); } else { assert_eq!(inp1.shape()[2], inp2.shape()[1]); } let out_shape = if adj_y { vec![inp1.shape()[0], inp1.shape()[1], inp2.shape()[1]] } else { vec![inp1.shape()[0], inp1.shape()[1], inp2.shape()[2]] }; let fc_chip = FullyConnectedChip:: { _marker: PhantomData, config: FullyConnectedConfig::construct(true), }; let mut outp: Vec> = vec![]; for i in 0..inp1.shape()[0] { let inp1_slice = inp1.index_axis(Axis(0), i).to_owned(); // Due to tensorflow BS, transpose the "weights" let inp2_slice = if adj_y { inp2.index_axis(Axis(0), i).to_owned() } else { inp2.index_axis(Axis(0), i).t().to_owned() }; println!("inp1_slice: {:?}", inp1_slice.shape()); println!("inp2_slice: {:?}", inp2_slice.shape()); // Batch MM doesn't have a fused activation, so insert it here // TODO: consider putting this in the converter? let tmp_config = LayerConfig { layer_params: vec![0], ..layer_config.clone() }; let outp_slice = fc_chip.forward( layouter.namespace(|| ""), &vec![inp1_slice, inp2_slice], constants, gadget_config.clone(), &tmp_config, )?; outp.extend(outp_slice[0].iter().map(|x| x.clone()).collect::>()); } let outp = Array::from_shape_vec(IxDyn(out_shape.as_slice()), outp).unwrap(); Ok(vec![outp]) } } impl GadgetConsumer for BatchMatMulChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::Adder, GadgetType::DotProduct, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/conv2d.rs ================================================ // TODO: Speed up Depthwise operations with Freivald's algorithm use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ bias_div_round_relu6::BiasDivRoundRelu6Chip, dot_prod::DotProductChip, gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::relu::ReluChip, }, layers::{ fully_connected::{FullyConnectedChip, FullyConnectedConfig}, shape::pad::pad, }, }; use super::layer::{ActivationType, AssignedTensor, GadgetConsumer, Layer, LayerConfig}; #[derive(Default, Clone, Copy, Eq, PartialEq)] pub enum PaddingEnum { #[default] Same, Valid, } #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] pub enum ConvLayerEnum { #[default] Conv2D, DepthwiseConv2D, } pub struct Conv2DConfig { pub conv_type: ConvLayerEnum, pub padding: PaddingEnum, pub activation: ActivationType, pub stride: (usize, usize), } pub struct Conv2DChip { pub config: LayerConfig, pub _marker: PhantomData, } impl Conv2DChip { // TODO: this is horrible. What's the best way to fix this? pub fn param_vec_to_config(layer_params: Vec) -> Conv2DConfig { let conv_type = match layer_params[0] { 0 => ConvLayerEnum::Conv2D, 1 => ConvLayerEnum::DepthwiseConv2D, _ => panic!("Invalid conv type"), }; let padding = match layer_params[1] { 0 => PaddingEnum::Same, 1 => PaddingEnum::Valid, _ => panic!("Invalid padding"), }; let activation = match layer_params[2] { 0 => ActivationType::None, 1 => ActivationType::Relu, 3 => ActivationType::Relu6, _ => panic!("Invalid activation type"), }; let stride = (layer_params[3] as usize, layer_params[4] as usize); Conv2DConfig { conv_type, padding, activation, stride, } } pub fn get_padding( h: usize, w: usize, si: usize, sj: usize, ci: usize, cj: usize, ) -> ((usize, usize), (usize, usize)) { let ph = if h % si == 0 { (ci as i64 - sj as i64).max(0) } else { (ci as i64 - (h % si) as i64).max(0) } as usize; let pw = if w % sj == 0 { (cj as i64 - sj as i64).max(0) } else { (cj as i64 - (w % sj) as i64).max(0) } as usize; ((ph / 2, ph - ph / 2), (pw / 2, pw - pw / 2)) } pub fn out_hw( h: usize, w: usize, si: usize, sj: usize, ch: usize, cw: usize, padding: PaddingEnum, ) -> (usize, usize) { /* println!( "H: {}, W: {}, SI: {}, SJ: {}, CH: {}, CW: {}", h, w, si, sj, ch, cw ); */ // https://iq.opengenus.org/same-and-valid-padding/ match padding { PaddingEnum::Same => ((h + si - 1) / si, (w + sj - 1) / sj), // TODO: the above is probably correct, but we always have valid paddings // PaddingEnum::Same => (h / si, w / sj), PaddingEnum::Valid => ((h - ch) / si + 1, (w - cw) / sj + 1), } } pub fn splat( &self, tensors: &Vec, IxDyn>>, zero: Rc, ) -> (Vec>>, Vec>>, Vec>) { // assert_eq!(tensors.len(), 3); assert!(tensors.len() <= 3); let conv_config = &Self::param_vec_to_config(self.config.layer_params.clone()); let inp = &tensors[0]; let weights = &tensors[1]; let zero_arr = Array::from_elem(IxDyn(&vec![1]), zero.clone()); let biases = if tensors.len() == 3 { &tensors[2] } else { &zero_arr }; let h: usize = inp.shape()[1]; let w: usize = inp.shape()[2]; let ch: usize = weights.shape()[1]; let cw: usize = weights.shape()[2]; let (si, sj) = conv_config.stride; // B, H, W, C assert_eq!(inp.shape().len(), 4); let (ph, pw) = if conv_config.padding == PaddingEnum::Same { Self::get_padding(h, w, si, sj, ch, cw) } else { ((0, 0), (0, 0)) }; // println!("Padding: {:?}", (ph, pw)); let padding = vec![[0, 0], [ph.0, ph.1], [pw.0, pw.1], [0, 0]]; let inp_pad = pad(&inp, padding, &zero); let (oh, ow) = Self::out_hw(h, w, si, sj, ch, cw, conv_config.padding); let mut inp_cells = vec![]; let mut weights_cells = vec![]; let mut biases_cells = vec![]; let mut input_row_idx = 0; let mut weight_row_idx = 0; // (output_channels x inp_channels * C_H * C_W) for chan_out in 0..weights.shape()[0] { weights_cells.push(vec![]); for ci in 0..weights.shape()[1] { for cj in 0..weights.shape()[2] { for ck in 0..weights.shape()[3] { weights_cells[weight_row_idx].push(weights[[chan_out, ci, cj, ck]].clone()); } } } weight_row_idx += 1; } // (O_H * O_W x inp_channels * C_H * C_W) for batch in 0..inp.shape()[0] { for i in 0..oh { for j in 0..ow { inp_cells.push(vec![]); for ci in 0..weights.shape()[1] { for cj in 0..weights.shape()[2] { for ck in 0..weights.shape()[3] { let idx_i = i * si + ci; let idx_j = j * sj + cj; inp_cells[input_row_idx].push(inp_pad[[batch, idx_i, idx_j, ck]].clone()); } } } input_row_idx += 1; } } } for _batch in 0..inp.shape()[0] { for _ in 0..oh { for _ in 0..ow { for chan_out in 0..weights.shape()[0] { if tensors.len() == 3 { biases_cells.push(biases[chan_out].clone()); } else { biases_cells.push(zero.clone()); } } } } } (inp_cells, weights_cells, biases_cells) } pub fn splat_depthwise( &self, tensors: &Vec, IxDyn>>, zero: Rc, ) -> (Vec>>, Vec>>, Vec>) { let input = &tensors[0]; let weights = &tensors[1]; let biases = &tensors[2]; assert_eq!(tensors.len(), 3); assert_eq!(input.shape().len(), 4); assert_eq!(weights.shape().len(), 4); assert_eq!(input.shape()[0], 1); let conv_config = &Self::param_vec_to_config(self.config.layer_params.clone()); let strides = conv_config.stride; let h: usize = input.shape()[1]; let w: usize = input.shape()[2]; let ch: usize = weights.shape()[1]; let cw: usize = weights.shape()[2]; let (si, sj) = conv_config.stride; let (oh, ow) = Self::out_hw(h, w, si, sj, ch, cw, conv_config.padding); let (ph, pw) = if conv_config.padding == PaddingEnum::Same { Self::get_padding(h, w, si, sj, ch, cw) } else { ((0, 0), (0, 0)) }; let padding = vec![[0, 0], [ph.0, ph.1], [pw.0, pw.1], [0, 0]]; let inp_pad = pad(&input, padding, &zero); let mut inp_cells = vec![]; let mut weight_cells = vec![]; let mut biases_cells = vec![]; let mut row_idx = 0; for i in 0..oh { for j in 0..ow { for chan_out in 0..weights.shape()[3] { inp_cells.push(vec![]); weight_cells.push(vec![]); biases_cells.push(biases[[chan_out]].clone()); for ci in 0..weights.shape()[1] { for cj in 0..weights.shape()[2] { let idx_i = i * strides.0 + ci; let idx_j = j * strides.1 + cj; inp_cells[row_idx].push(inp_pad[[0, idx_i, idx_j, chan_out]].clone()); weight_cells[row_idx].push(weights[[0, ci, cj, chan_out]].clone()); } } row_idx += 1; } } } (inp_cells, weight_cells, biases_cells) } } impl Layer for Conv2DChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let conv_config = &Self::param_vec_to_config(self.config.layer_params.clone()); let zero = constants.get(&0).unwrap(); let inp = &tensors[0]; let weights = &tensors[1]; let (oh, ow) = Self::out_hw( inp.shape()[1], inp.shape()[2], conv_config.stride.0, conv_config.stride.1, weights.shape()[1], weights.shape()[2], conv_config.padding, ); let batch_size = inp.shape()[0]; let (splat_inp, splat_weights, splat_biases) = match conv_config.conv_type { ConvLayerEnum::Conv2D => self.splat(tensors, zero.clone()), ConvLayerEnum::DepthwiseConv2D => self.splat_depthwise(tensors, zero.clone()), }; let outp_flat: Vec> = match conv_config.conv_type { ConvLayerEnum::Conv2D => { let fc_chip = FullyConnectedChip:: { _marker: PhantomData, config: FullyConnectedConfig::construct(false), }; let conv_size = splat_inp[0].len(); let flattened_inp: Vec<_> = splat_inp.into_iter().flat_map(|x| x.into_iter()).collect(); let flattened_weights = splat_weights .into_iter() .flat_map(|x| x.into_iter()) .collect::>(); let out_channels = weights.shape()[0]; let inp_array = Array::from_shape_vec(IxDyn(&vec![batch_size * oh * ow, conv_size]), flattened_inp) .unwrap(); let weights_array = Array::from_shape_vec(IxDyn(&vec![out_channels, conv_size]), flattened_weights).unwrap(); let outp_slice = fc_chip .forward( layouter.namespace(|| ""), &vec![weights_array, inp_array], constants, gadget_config.clone(), layer_config, ) .unwrap(); let outp_flat = outp_slice[0] .t() .into_iter() .map(|x| (**x).clone()) .collect::>(); outp_flat } ConvLayerEnum::DepthwiseConv2D => { // Do the dot products let dot_prod_chip = DotProductChip::::construct(gadget_config.clone()); let mut outp_flat = vec![]; for (inp_vec, weight_vec) in splat_inp.iter().zip(splat_weights.iter()) { let inp_vec = inp_vec.iter().map(|x| x.as_ref()).collect::>(); let weight_vec = weight_vec.iter().map(|x| x.as_ref()).collect::>(); let vec_inputs = vec![inp_vec, weight_vec]; let constants = vec![zero.as_ref()]; let outp = dot_prod_chip .forward(layouter.namespace(|| "dot_prod"), &vec_inputs, &constants) .unwrap(); outp_flat.push(outp[0].clone()); } // println!("outp_flat: {:?}", outp_flat.len()); outp_flat } }; let mut biases = vec![]; for bias in splat_biases.iter() { biases.push(bias.as_ref()); } // Compute the bias + div + relu let bdr_chip = BiasDivRoundRelu6Chip::::construct(gadget_config.clone()); let tmp = vec![zero.as_ref()]; let outp_flat = outp_flat.iter().map(|x| x).collect::>(); let outp = bdr_chip .forward( layouter.namespace(|| "bias_div_relu"), &vec![outp_flat, biases], &tmp, ) .unwrap(); // TODO: this is also horrible. The bdr chip outputs interleaved [(relu'd, div'd), (relu'd, div'd), ...] // Uninterleave depending on whether or not we're doing the relu let outp = if conv_config.activation == ActivationType::Relu6 { outp .into_iter() .step_by(2) .map(|x| Rc::new(x)) .collect::>() } else if conv_config.activation == ActivationType::None { outp .into_iter() .skip(1) .step_by(2) .map(|x| Rc::new(x)) .collect::>() } else if conv_config.activation == ActivationType::Relu { let dived = outp.iter().skip(1).step_by(2).collect::>(); let relu_chip = ReluChip::::construct(gadget_config.clone()); let relu_outp = relu_chip .forward(layouter.namespace(|| "relu"), &vec![dived], &tmp) .unwrap(); let relu_outp = relu_outp .into_iter() .map(|x| Rc::new(x)) .collect::>(); relu_outp } else { panic!("Unsupported activation type"); }; let oc = match conv_config.conv_type { ConvLayerEnum::Conv2D => weights.shape()[0], ConvLayerEnum::DepthwiseConv2D => weights.shape()[3], }; let out_shape = vec![batch_size, oh, ow, oc]; let outp = Array::from_shape_vec(IxDyn(&out_shape), outp).unwrap(); Ok(vec![outp]) } } impl GadgetConsumer for Conv2DChip { fn used_gadgets(&self, layer_params: Vec) -> Vec { let conv_config = &Self::param_vec_to_config(layer_params.clone()); let mut outp = vec![ GadgetType::Adder, GadgetType::DotProduct, GadgetType::InputLookup, GadgetType::BiasDivRoundRelu6, ]; if conv_config.activation == ActivationType::Relu { outp.push(GadgetType::Relu); } outp } } ================================================ FILE: src/layers/dag.rs ================================================ use std::{collections::HashMap, fs::File, io::BufWriter, marker::PhantomData, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use crate::{ gadgets::gadget::{convert_to_u64, GadgetConfig}, layers::{ arithmetic::{add::AddChip, div_var::DivVarChip, mul::MulChip, sub::SubChip}, batch_mat_mul::BatchMatMulChip, div_fixed::DivFixedChip, fully_connected::{FullyConnectedChip, FullyConnectedConfig}, logistic::LogisticChip, max_pool_2d::MaxPool2DChip, mean::MeanChip, noop::NoopChip, pow::PowChip, rsqrt::RsqrtChip, shape::{ broadcast::BroadcastChip, concatenation::ConcatenationChip, mask_neg_inf::MaskNegInfChip, pack::PackChip, pad::PadChip, permute::PermuteChip, reshape::ReshapeChip, resize_nn::ResizeNNChip, rotate::RotateChip, slice::SliceChip, split::SplitChip, transpose::TransposeChip, }, softmax::SoftmaxChip, sqrt::SqrtChip, square::SquareChip, squared_diff::SquaredDiffChip, tanh::TanhChip, update::UpdateChip, }, utils::helpers::print_assigned_arr, }; use super::{ avg_pool_2d::AvgPool2DChip, conv2d::Conv2DChip, layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig, LayerType}, }; #[derive(Clone, Debug, Default)] pub struct DAGLayerConfig { pub ops: Vec, pub inp_idxes: Vec>, pub out_idxes: Vec>, pub final_out_idxes: Vec, } pub struct DAGLayerChip { dag_config: DAGLayerConfig, _marker: PhantomData, } impl DAGLayerChip { pub fn construct(dag_config: DAGLayerConfig) -> Self { Self { dag_config, _marker: PhantomData, } } // IMPORTANT: Assumes input tensors are in order. Output tensors can be in any order. pub fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result<(HashMap>, Vec>), Error> { // Tensor map let mut tensor_map = HashMap::new(); for (idx, tensor) in tensors.iter().enumerate() { tensor_map.insert(idx, tensor.clone()); } // Compute the dag for (layer_idx, layer_config) in self.dag_config.ops.iter().enumerate() { let layer_type = &layer_config.layer_type; let inp_idxes = &self.dag_config.inp_idxes[layer_idx]; let out_idxes = &self.dag_config.out_idxes[layer_idx]; println!( "Processing layer {}, type: {:?}, inp_idxes: {:?}, out_idxes: {:?}, layer_params: {:?}", layer_idx, layer_type, inp_idxes, out_idxes, layer_config.layer_params ); let vec_inps = inp_idxes .iter() .map(|idx| tensor_map.get(idx).unwrap().clone()) .collect::>(); let out = match layer_type { LayerType::Add => { let add_chip = AddChip {}; add_chip.forward( layouter.namespace(|| "dag add"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::AvgPool2D => { let avg_pool_2d_chip = AvgPool2DChip {}; avg_pool_2d_chip.forward( layouter.namespace(|| "dag avg pool 2d"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::MaxPool2D => { let max_pool_2d_chip = MaxPool2DChip { marker: PhantomData::, }; max_pool_2d_chip.forward( layouter.namespace(|| "dag max pool 2d"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::BatchMatMul => { let batch_mat_mul_chip = BatchMatMulChip {}; batch_mat_mul_chip.forward( layouter.namespace(|| "dag batch mat mul"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Broadcast => { let broadcast_chip = BroadcastChip {}; broadcast_chip.forward( layouter.namespace(|| "dag batch mat mul"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Conv2D => { let conv_2d_chip = Conv2DChip { config: layer_config.clone(), _marker: PhantomData, }; conv_2d_chip.forward( layouter.namespace(|| "dag conv 2d"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::DivFixed => { let div_fixed_chip = DivFixedChip {}; div_fixed_chip.forward( layouter.namespace(|| "dag div"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::DivVar => { let div_var_chip = DivVarChip {}; div_var_chip.forward( layouter.namespace(|| "dag div"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::FullyConnected => { let fc_chip = FullyConnectedChip { _marker: PhantomData, config: FullyConnectedConfig::construct(true), }; fc_chip.forward( layouter.namespace(|| "dag fully connected"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Softmax => { let softmax_chip = SoftmaxChip {}; softmax_chip.forward( layouter.namespace(|| "dag softmax"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Mean => { let mean_chip = MeanChip {}; mean_chip.forward( layouter.namespace(|| "dag mean"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Pad => { let pad_chip = PadChip {}; pad_chip.forward( layouter.namespace(|| "dag pad"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Permute => { let pad_chip = PermuteChip {}; pad_chip.forward( layouter.namespace(|| "dag permute"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::SquaredDifference => { let squared_diff_chip = SquaredDiffChip {}; squared_diff_chip.forward( layouter.namespace(|| "dag squared diff"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Rsqrt => { let rsqrt_chip = RsqrtChip {}; rsqrt_chip.forward( layouter.namespace(|| "dag rsqrt"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Sqrt => { let sqrt_chip = SqrtChip {}; sqrt_chip.forward( layouter.namespace(|| "dag sqrt"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Logistic => { let logistic_chip = LogisticChip {}; logistic_chip.forward( layouter.namespace(|| "dag logistic"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Pow => { let pow_chip = PowChip {}; pow_chip.forward( layouter.namespace(|| "dag logistic"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Tanh => { let tanh_chip = TanhChip {}; tanh_chip.forward( layouter.namespace(|| "dag tanh"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Mul => { let mul_chip = MulChip {}; mul_chip.forward( layouter.namespace(|| "dag mul"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Sub => { let sub_chip = SubChip {}; sub_chip.forward( layouter.namespace(|| "dag sub"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Noop => { let noop_chip = NoopChip {}; noop_chip.forward( layouter.namespace(|| "dag noop"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Transpose => { let transpose_chip = TransposeChip {}; transpose_chip.forward( layouter.namespace(|| "dag transpose"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Reshape => { let reshape_chip = ReshapeChip {}; reshape_chip.forward( layouter.namespace(|| "dag reshape"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::ResizeNN => { let resize_nn_chip = ResizeNNChip {}; resize_nn_chip.forward( layouter.namespace(|| "dag resize nn"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Rotate => { let rotate_chip = RotateChip {}; rotate_chip.forward( layouter.namespace(|| "dag rotate"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Concatenation => { let concat_chip = ConcatenationChip {}; concat_chip.forward( layouter.namespace(|| "dag concatenation"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Pack => { let pack_chip = PackChip {}; pack_chip.forward( layouter.namespace(|| "dag pack"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Split => { let split_chip = SplitChip {}; split_chip.forward( layouter.namespace(|| "dag split"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Update => { let split_chip = UpdateChip {}; split_chip.forward( layouter.namespace(|| "dag update"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Slice => { let slice_chip = SliceChip {}; slice_chip.forward( layouter.namespace(|| "dag slice"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::MaskNegInf => { let mask_neg_inf_chip = MaskNegInfChip {}; mask_neg_inf_chip.forward( layouter.namespace(|| "dag mask neg inf"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } LayerType::Square => { let square_chip = SquareChip {}; square_chip.forward( layouter.namespace(|| "dag square"), &vec_inps, constants, gadget_config.clone(), &layer_config, )? } }; for (idx, tensor_idx) in out_idxes.iter().enumerate() { println!("Out {} shape: {:?}", idx, out[idx].shape()); tensor_map.insert(*tensor_idx, out[idx].clone()); } println!(); } let mut final_out = vec![]; for idx in self.dag_config.final_out_idxes.iter() { final_out.push(tensor_map.get(idx).unwrap().clone()); } let print_arr = if final_out.len() > 0 { &final_out[0] } else { if self.dag_config.ops.len() > 0 { let last_layer_idx = self.dag_config.ops.len() - 1; let out_idx = self.dag_config.out_idxes[last_layer_idx][0]; tensor_map.get(&out_idx).unwrap() } else { tensor_map.get(&0).unwrap() } }; let tmp = print_arr.iter().map(|x| x.as_ref()).collect::>(); print_assigned_arr("final out", &tmp.to_vec(), gadget_config.scale_factor); println!("final out idxes: {:?}", self.dag_config.final_out_idxes); let mut x = vec![]; for cell in print_arr.iter() { cell.value().map(|v| { let bias = 1 << 60 as i64; let v_pos = *v + F::from(bias as u64); let v = convert_to_u64(&v_pos) as i64 - bias; x.push(v); }); } if x.len() > 0 { let out_fname = "out.msgpack"; let f = File::create(out_fname).unwrap(); let mut buf = BufWriter::new(f); rmp_serde::encode::write_named(&mut buf, &x).unwrap(); } Ok((tensor_map, final_out)) } } impl GadgetConsumer for DAGLayerChip { // Special case: DAG doesn't do anything fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/div_fixed.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Value}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, var_div::VarDivRoundChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct DivFixedChip {} impl DivFixedChip { fn get_div_val( &self, mut layouter: impl Layouter, _tensors: &Vec>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result, Error> { // FIXME: this needs to be revealed let div = layer_config.layer_params[0]; let div = F::from(div as u64); let div = layouter .assign_region( || "division", |mut region| { let div = region .assign_advice( || "avg pool 2d div", gadget_config.columns[0], 0, || Value::known(div), ) .unwrap(); Ok(div) }, ) .unwrap(); Ok(div) } } impl Layer for DivFixedChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let inp_flat = inp.iter().map(|x| x.as_ref()).collect::>(); let zero = constants.get(&0).unwrap().as_ref(); let shape = inp.shape(); let div = self.get_div_val( layouter.namespace(|| "average div"), tensors, gadget_config.clone(), layer_config, )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let dived = var_div_chip.forward( layouter.namespace(|| "average div"), &vec![inp_flat], &vec![zero, &div], )?; let dived = dived.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(shape), dived).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for DivFixedChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::VarDivRound] } } ================================================ FILE: src/layers/fully_connected.rs ================================================ use std::{collections::HashMap, marker::PhantomData, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Region, Value}, halo2curves::ff::PrimeField, plonk::{Advice, Column, Error}, }; use ndarray::{Array, ArrayView, Axis, IxDyn}; use crate::{ gadgets::{ add_pairs::AddPairsChip, dot_prod::DotProductChip, gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::relu::ReluChip, var_div::VarDivRoundChip, }, layers::layer::ActivationType, utils::helpers::RAND_START_IDX, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; pub struct FullyConnectedConfig { pub normalize: bool, // Should be true } impl FullyConnectedConfig { pub fn construct(normalize: bool) -> Self { Self { normalize } } } pub struct FullyConnectedChip { pub _marker: PhantomData, pub config: FullyConnectedConfig, } impl FullyConnectedChip { pub fn compute_mm( // input: &AssignedTensor, input: &ArrayView, IxDyn>, weight: &AssignedTensor, ) -> Array, IxDyn> { assert_eq!(input.ndim(), 2); assert_eq!(weight.ndim(), 2); assert_eq!(input.shape()[1], weight.shape()[0]); let mut outp = vec![]; for i in 0..input.shape()[0] { for j in 0..weight.shape()[1] { let mut sum = input[[i, 0]].value().map(|x: &F| *x) * weight[[0, j]].value(); for k in 1..input.shape()[1] { sum = sum + input[[i, k]].value().map(|x: &F| *x) * weight[[k, j]].value(); } outp.push(sum); } } let out_shape = [input.shape()[0], weight.shape()[1]]; Array::from_shape_vec(IxDyn(out_shape.as_slice()), outp).unwrap() } pub fn assign_array( columns: &Vec>, region: &mut Region, array: &Array, IxDyn>, ) -> Result, IxDyn>, Error> { assert_eq!(array.ndim(), 2); let mut outp = vec![]; for (idx, val) in array.iter().enumerate() { let row_idx = idx / columns.len(); let col_idx = idx % columns.len(); let cell = region .assign_advice(|| "assign array", columns[col_idx], row_idx, || *val) .unwrap(); outp.push(cell); } let out_shape = [array.shape()[0], array.shape()[1]]; Ok(Array::from_shape_vec(IxDyn(out_shape.as_slice()), outp).unwrap()) } pub fn random_vector( constants: &HashMap>, size: usize, ) -> Result>, Error> { let mut outp = vec![]; for idx in 0..size { let idx = RAND_START_IDX + (idx as i64); if !constants.contains_key(&idx) { println!("Random vector is too small: {:?}", size); } let cell = constants.get(&idx).unwrap().clone(); outp.push(cell); } Ok(outp) } fn get_activation(&self, layer_params: &Vec) -> ActivationType { let activation = layer_params[0]; match activation { 0 => ActivationType::None, 1 => ActivationType::Relu, _ => panic!("Unsupported activation type for fully connected"), } } } impl Layer for FullyConnectedChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { assert!(tensors.len() <= 3); let activation = self.get_activation(&layer_config.layer_params); let input = &tensors[0]; let ndim = input.ndim(); let input = if ndim == 2 { ArrayView::from(input) } else { input.index_axis(Axis(0), 0) }; let weight = &tensors[1].t().into_owned(); let zero = constants.get(&0).unwrap().as_ref(); // Compute and assign the result let mm_result = layouter .assign_region( || "compute and assign mm", |mut region| { let mm_result = Self::compute_mm(&input, weight); let mm_result = Self::assign_array(&gadget_config.columns, &mut region, &mm_result).unwrap(); Ok(mm_result) }, ) .unwrap(); // Generate random vectors let r1 = Self::random_vector(constants, mm_result.shape()[0]).unwrap(); let r2 = Self::random_vector(constants, mm_result.shape()[1]).unwrap(); let dot_prod_chip = DotProductChip::::construct(gadget_config.clone()); let r1_ref = r1.iter().map(|x| x.as_ref()).collect::>(); let r2_ref = r2.iter().map(|x| x.as_ref()).collect::>(); // Compute r1 * result let mut r1_res = vec![]; // println!("r1_ref: {:?}", r1_ref.len()); // println!("r2_ref: {:?}", r2_ref.len()); // println!("mm_result: {:?}", mm_result.shape()); for i in 0..mm_result.shape()[1] { let tmp = mm_result.index_axis(Axis(1), i); let mm_ci = tmp.iter().collect::>(); let r1_res_i = dot_prod_chip .forward( layouter.namespace(|| format!("r1_res_{}", i)), &vec![mm_ci, r1_ref.clone()], &vec![zero], ) .unwrap(); r1_res.push(r1_res_i[0].clone()); } // Compute r1 * result * r2 let r1_res_ref = r1_res.iter().collect::>(); let r1_res_r2 = dot_prod_chip .forward( layouter.namespace(|| "r1_res_r2"), &vec![r1_res_ref, r2_ref.clone()], &vec![zero], ) .unwrap(); let r1_res_r2 = r1_res_r2[0].clone(); // println!("r1_res_r2: {:?}", r1_res_r2); // Compute r1 * input let mut r1_input = vec![]; // println!("input: {:?}", input.shape()); // println!("r1_ref: {:?}", r1_ref.len()); for i in 0..input.shape()[1] { let tmp = input.index_axis(Axis(1), i); let input_ci = tmp.iter().map(|x| x.as_ref()).collect::>(); let r1_input_i = dot_prod_chip .forward( layouter.namespace(|| format!("r1_input_{}", i)), &vec![input_ci, r1_ref.clone()], &vec![zero], ) .unwrap(); r1_input.push(r1_input_i[0].clone()); } // Compute weight * r2 let mut weight_r2 = vec![]; for i in 0..weight.shape()[0] { let tmp = weight.index_axis(Axis(0), i); let weight_ci = tmp.iter().map(|x| x.as_ref()).collect::>(); let weight_r2_i = dot_prod_chip .forward( layouter.namespace(|| format!("weight_r2_{}", i)), &vec![weight_ci, r2_ref.clone()], &vec![zero], ) .unwrap(); weight_r2.push(weight_r2_i[0].clone()); } // Compute (r1 * input) * (weight * r2) let r1_input_ref = r1_input.iter().collect::>(); let weight_r2_ref = weight_r2.iter().collect::>(); let r1_inp_weight_r2 = dot_prod_chip .forward( layouter.namespace(|| "r1_inp_weight_r2"), &vec![r1_input_ref, weight_r2_ref], &vec![zero], ) .unwrap(); let r1_inp_weight_r2 = r1_inp_weight_r2[0].clone(); // println!("r1_inp_weight_r2: {:?}", r1_inp_weight_r2); layouter .assign_region( || "fc equality check", |mut region| { let t1 = r1_res_r2 .copy_advice(|| "", &mut region, gadget_config.columns[0], 0) .unwrap(); let t2 = r1_inp_weight_r2 .copy_advice(|| "", &mut region, gadget_config.columns[0], 1) .unwrap(); region.constrain_equal(t1.cell(), t2.cell()).unwrap(); Ok(()) }, ) .unwrap(); let shape = [mm_result.shape()[0], mm_result.shape()[1]]; let final_result_flat = if self.config.normalize { let mm_flat = mm_result.iter().collect::>(); let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let sf = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); let mm_div = var_div_chip .forward( layouter.namespace(|| "mm_div"), &vec![mm_flat], &vec![zero, sf], ) .unwrap(); let mm_div = if tensors.len() == 3 { let bias = tensors[2].broadcast(shape.clone()).unwrap(); let bias = bias.iter().map(|x| x.as_ref()).collect::>(); let mm_div = mm_div.iter().collect::>(); let adder_chip = AddPairsChip::::construct(gadget_config.clone()); let mm_bias = adder_chip .forward( layouter.namespace(|| "mm_bias"), &vec![mm_div, bias], &vec![zero], ) .unwrap(); mm_bias } else { mm_div }; let mm_div = if activation == ActivationType::Relu { let relu_chip = ReluChip::::construct(gadget_config.clone()); let mm_div = mm_div.iter().collect::>(); let vec_inputs = vec![mm_div]; relu_chip .forward(layouter.namespace(|| "relu"), &vec_inputs, &vec![zero]) .unwrap() } else if activation == ActivationType::None { mm_div } else { panic!("Unsupported activation type"); }; mm_div.into_iter().map(|x| Rc::new(x)).collect::>() } else { mm_result .into_iter() .map(|x| Rc::new(x)) .collect::>() }; let final_result = Array::from_shape_vec(IxDyn(&shape), final_result_flat).unwrap(); Ok(vec![final_result]) } } impl GadgetConsumer for FullyConnectedChip { fn used_gadgets(&self, layer_params: Vec) -> Vec { let activation = self.get_activation(&layer_params); let mut outp = vec![ GadgetType::Adder, GadgetType::AddPairs, GadgetType::DotProduct, GadgetType::VarDivRound, GadgetType::InputLookup, ]; match activation { ActivationType::Relu => outp.push(GadgetType::Relu), ActivationType::None => (), _ => panic!("Unsupported activation type"), } outp } } ================================================ FILE: src/layers/layer.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, IxDyn}; use crate::gadgets::gadget::{GadgetConfig, GadgetType}; #[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq)] pub enum LayerType { Add, AvgPool2D, BatchMatMul, Broadcast, Concatenation, Conv2D, DivVar, DivFixed, FullyConnected, Logistic, MaskNegInf, MaxPool2D, Mean, Mul, #[default] Noop, Pack, Pad, Pow, Permute, Reshape, ResizeNN, Rotate, Rsqrt, Slice, Softmax, Split, Sqrt, Square, SquaredDifference, Sub, Tanh, Transpose, Update, } // NOTE: This is the same order as the TFLite schema // Must not be changed #[derive(Clone, Debug, Default, Hash, Eq, PartialEq)] pub enum ActivationType { #[default] None, Relu, ReluN1To1, Relu6, Tanh, SignBit, } #[derive(Clone, Debug, Default)] pub struct LayerConfig { pub layer_type: LayerType, pub layer_params: Vec, // This is turned into layer specific configurations at runtime pub inp_shapes: Vec>, pub out_shapes: Vec>, pub mask: Vec, } pub type CellRc = Rc>; pub type AssignedTensor = Array, IxDyn>; // General issue with rust: I'm not sure how to pass named arguments to a trait... // Currently, the caller must be aware of the order of the tensors and results pub trait Layer { fn forward( &self, layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error>; } pub trait GadgetConsumer { fn used_gadgets(&self, layer_params: Vec) -> Vec; } ================================================ FILE: src/layers/logistic.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::logistic::LogisticGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct LogisticChip {} impl Layer for LogisticChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let inp_vec = inp.iter().map(|x| x.as_ref()).collect::>(); let zero = constants.get(&0).unwrap().as_ref(); let logistic_chip = LogisticGadgetChip::::construct(gadget_config.clone()); let vec_inps = vec![inp_vec]; let constants = vec![zero]; let out = logistic_chip.forward( layouter.namespace(|| "logistic chip"), &vec_inps, &constants, )?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for LogisticChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Logistic, GadgetType::InputLookup] } } ================================================ FILE: src/layers/max_pool_2d.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, max::MaxChip, }, layers::conv2d::{Conv2DChip, PaddingEnum}, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; pub struct MaxPool2DChip { pub marker: std::marker::PhantomData, } impl MaxPool2DChip { pub fn shape(inp: &AssignedTensor, layer_config: &LayerConfig) -> (usize, usize) { let params = &layer_config.layer_params; let (fx, fy) = (params[0], params[1]); let (fx, fy) = (fx as usize, fy as usize); let (sx, sy) = (params[2], params[3]); let (sx, sy) = (sx as usize, sy as usize); // Only support batch size 1 for now assert_eq!(inp.shape()[0], 1); let out_shape = Conv2DChip::::out_hw( inp.shape()[1], inp.shape()[2], sx, sy, fx, fy, PaddingEnum::Valid, ); out_shape } pub fn splat( inp: &AssignedTensor, layer_config: &LayerConfig, ) -> Result>>, Error> { let params = &layer_config.layer_params; let (fx, fy) = (params[0], params[1]); let (fx, fy) = (fx as usize, fy as usize); let (sx, sy) = (params[2], params[3]); let (sx, sy) = (sx as usize, sy as usize); // Only support batch size 1 for now assert_eq!(inp.shape()[0], 1); let out_shape = Self::shape(inp, layer_config); let mut splat = vec![]; for i in 0..out_shape.0 { for j in 0..out_shape.1 { for k in 0..inp.shape()[3] { let mut tmp = vec![]; for x in 0..fx { for y in 0..fy { let x = i * sx + x; let y = j * sy + y; if x < inp.shape()[1] && y < inp.shape()[2] { tmp.push(inp[[0, x, y, k]].clone()); } } } splat.push(tmp); } } } Ok(splat) } } impl Layer for MaxPool2DChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let splat = Self::splat(inp, layer_config).unwrap(); let max_chip = MaxChip::::construct(gadget_config.clone()); let mut out = vec![]; for i in 0..splat.len() { let inps = &splat[i]; let inps = inps.iter().map(|x| x.as_ref()).collect(); let max = max_chip .forward( layouter.namespace(|| format!("max {}", i)), &vec![inps], &vec![], ) .unwrap(); out.push(max[0].clone()); } let out = out.into_iter().map(|x| Rc::new(x)).collect(); // TODO: refactor this let out_xy = Self::shape(inp, layer_config); let out_shape = vec![1, out_xy.0, out_xy.1, inp.shape()[3]]; let out = Array::from_shape_vec(IxDyn(&out_shape), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for MaxPool2DChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Max, GadgetType::InputLookup] } } ================================================ FILE: src/layers/mean.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter, Value}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, Axis, IxDyn}; use crate::gadgets::gadget::{GadgetConfig, GadgetType}; use super::{ averager::Averager, layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}, }; pub struct MeanChip {} impl MeanChip { pub fn get_keep_axis(&self, layer_config: &LayerConfig) -> usize { let inp_shape = &layer_config.inp_shapes[0]; let out_shape = &layer_config.out_shapes[0]; assert_eq!(inp_shape[0], 1); assert_eq!(out_shape[0], 1); // Skip the batch axis let mut keep_axes = (1..inp_shape.len()).collect::>(); for mean_axis in layer_config.layer_params.iter() { keep_axes.retain(|&x| x != *mean_axis as usize); } assert_eq!(keep_axes.len(), 1); keep_axes[0] /* let mut num_same = 0; let mut keep_axis: i64 = -1; for i in 1..inp_shape.len() { if inp_shape[i] == out_shape[i] { keep_axis = i as i64; num_same += 1; } } if keep_axis == -1 { panic!("All axes are different"); } if num_same > 1 { panic!("More than one axis is the same"); } keep_axis as usize */ } } impl Averager for MeanChip { fn splat(&self, input: &AssignedTensor, layer_config: &LayerConfig) -> Vec>> { // Only support batch size = 1 assert_eq!(input.shape()[0], 1); // Only support batch + 2D, summing over one axis // assert_eq!(input.shape().len(), 3); let keep_axis = self.get_keep_axis(layer_config); let mut splat = vec![]; for i in 0..input.shape()[keep_axis] { let mut tmp = vec![]; for x in input.index_axis(Axis(keep_axis), i).iter() { tmp.push(x.clone()); } splat.push(tmp); } splat } fn get_div_val( &self, mut layouter: impl Layouter, tensors: &Vec>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result, Error> { let inp = &tensors[0]; let keep_axis = self.get_keep_axis(layer_config); let mut div = 1; for i in 0..inp.shape().len() { if i != keep_axis { div *= inp.shape()[i]; } } let div = F::from(div as u64); // FIXME: put this in the fixed column let div = layouter.assign_region( || "mean div", |mut region| { let div = region.assign_advice( || "mean div", gadget_config.columns[0], 0, || Value::known(div), )?; Ok(div) }, )?; Ok(div) } } impl Layer for MeanChip { fn forward( &self, layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let dived = self.avg_forward(layouter, tensors, constants, gadget_config, layer_config)?; let out_shape = layer_config.out_shapes[0] .iter() .map(|x| *x as usize) .collect::>(); let out = Array::from_shape_vec(IxDyn(&out_shape), dived).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for MeanChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::Adder, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/noop.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use crate::gadgets::gadget::GadgetConfig; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; pub struct NoopChip {} impl Layer for NoopChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let ret_idx = layer_config.layer_params[0] as usize; Ok(vec![tensors[ret_idx].clone()]) } } impl GadgetConsumer for NoopChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/pow.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::pow::PowGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct PowChip {} impl Layer for PowChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let inp_vec = inp.iter().map(|x| x.as_ref()).collect::>(); let zero = constants.get(&0).unwrap().as_ref(); let pow_chip = PowGadgetChip::::construct(gadget_config.clone()); let vec_inps = vec![inp_vec]; let constants = vec![zero]; let out = pow_chip.forward(layouter.namespace(|| "pow chip"), &vec_inps, &constants)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for PowChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Pow, GadgetType::InputLookup] } } ================================================ FILE: src/layers/rsqrt.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::rsqrt::RsqrtGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct RsqrtChip {} impl Layer for RsqrtChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let mut inp_vec = vec![]; let mask = &layer_config.mask; let mut mask_map = HashMap::new(); for i in 0..mask.len() / 2 { mask_map.insert(mask[2 * i], mask[2 * i + 1]); } let min_val = gadget_config.min_val; let min_val = constants.get(&min_val).unwrap().as_ref(); let max_val = gadget_config.max_val; let max_val = constants.get(&max_val).unwrap().as_ref(); for (i, val) in inp.iter().enumerate() { let i = i as i64; if mask_map.contains_key(&i) { let mask_val = *mask_map.get(&i).unwrap(); if mask_val == 1 { inp_vec.push(max_val); } else if mask_val == -1 { inp_vec.push(min_val); } else { panic!(); } } else { inp_vec.push(val.as_ref()); } } let zero = constants.get(&0).unwrap().as_ref(); let rsqrt_chip = RsqrtGadgetChip::::construct(gadget_config.clone()); let vec_inps = vec![inp_vec]; let constants = vec![zero, min_val, max_val]; let out = rsqrt_chip.forward(layouter.namespace(|| "rsqrt chip"), &vec_inps, &constants)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for RsqrtChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Rsqrt, GadgetType::InputLookup] } } ================================================ FILE: src/layers/shape/broadcast.rs ================================================ // // Broadcast is used as a temporary measure to represent a the backprop // of a full-kernel AvgPool2D // use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::Array; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct BroadcastChip {} // TODO: Fix this after demo impl Layer for BroadcastChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let shape = inp.shape(); let output_shape = layer_config.out_shapes[0].clone(); // Check that we only broadcast dimensions with shape 1 assert!(shape.len() == output_shape.len()); assert!(shape.len() == 4); for (inp, outp) in shape.iter().zip(output_shape.iter()) { if *inp != *outp && !(*inp == 1) { panic!(); } } let mut output_flat = vec![]; for i in 0..output_shape[0] { for j in 0..output_shape[1] { for k in 0..output_shape[2] { for l in 0..output_shape[3] { let indexes = [i, j, k, l] .iter() .enumerate() .map(|(idx, x)| if shape[idx] == 1 { 0 } else { *x }) .collect::>(); output_flat.push(inp[[indexes[0], indexes[1], indexes[2], indexes[3]]].clone()); } } } } println!("Broadcast : {:?} -> {:?}", inp.shape(), output_shape); let out = Array::from_shape_vec(output_shape, output_flat).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for BroadcastChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/concatenation.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{concatenate, Axis}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct ConcatenationChip {} impl Layer for ConcatenationChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let axis = layer_config.layer_params[0] as usize; let views = tensors.iter().map(|x| x.view()).collect::>(); // TODO: this is a bit of a hack let out = concatenate(Axis(axis), views.as_slice()).unwrap_or(tensors[0].clone()); Ok(vec![out]) } } impl GadgetConsumer for ConcatenationChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/mask_neg_inf.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct MaskNegInfChip {} impl Layer for MaskNegInfChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let mask_ndim = layer_config.layer_params[0] as usize; let mask_shape = layer_config.layer_params[1..mask_ndim + 1] .iter() .map(|x| *x as usize) .collect::>(); let mask_vec = layer_config.layer_params[mask_ndim + 1..].to_vec(); let mask = Array::from_shape_vec(IxDyn(&mask_shape), mask_vec).unwrap(); let mask = mask.broadcast(inp.raw_dim()).unwrap(); let min_val = gadget_config.min_val; let min_val = constants.get(&min_val).unwrap().clone(); let mut out_vec = vec![]; for (val, to_mask) in inp.iter().zip(mask.iter()) { if *to_mask == 0 { out_vec.push(val.clone()); } else { out_vec.push(min_val.clone()); } } let outp = Array::from_shape_vec(inp.raw_dim(), out_vec).unwrap(); Ok(vec![outp]) } } impl GadgetConsumer for MaskNegInfChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/pack.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{concatenate, Axis}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct PackChip {} impl Layer for PackChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let axis = layer_config.layer_params[0] as usize; if axis > 1 { panic!("Pack only supports axis=0 or axis=1"); } let expanded = tensors .into_iter() .map(|x| x.clone().insert_axis(Axis(axis))) .collect::>(); let views = expanded.iter().map(|x| x.view()).collect::>(); // TODO: in some cases, the pack is unnecessary. Simply return the first tensor in this case let out = concatenate(Axis(axis), views.as_slice()).unwrap_or(tensors[0].clone()); Ok(vec![out]) } } impl GadgetConsumer for PackChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/pad.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{Array, Axis, IxDyn, Slice}; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; // TODO: figure out where to put this pub fn pad( input: &Array, IxDyn>, padding: Vec<[usize; 2]>, pad_val: &Rc, ) -> Array, IxDyn> { let tmp = input.iter().collect(); let input = Array::from_shape_vec(input.raw_dim(), tmp).unwrap(); assert_eq!(input.ndim(), padding.len()); let mut padded_shape = input.raw_dim(); for (ax, (&ax_len, &[pad_lo, pad_hi])) in input.shape().iter().zip(&padding).enumerate() { padded_shape[ax] = ax_len + pad_lo + pad_hi; } let mut padded = Array::from_elem(padded_shape, pad_val); let padded_dim = padded.raw_dim(); { // Select portion of padded array that needs to be copied from the // original array. let mut orig_portion = padded.view_mut(); for (ax, &[pad_lo, pad_hi]) in padding.iter().enumerate() { orig_portion.slice_axis_inplace( Axis(ax), Slice::from(pad_lo as isize..padded_dim[ax] as isize - (pad_hi as isize)), ); } // Copy the data from the original array. orig_portion.assign(&input.view()); } let dim = padded.raw_dim(); let tmp = padded.into_iter().map(|x| x.clone()).collect(); let padded = Array::from_shape_vec(dim, tmp).unwrap(); padded } pub struct PadChip {} pub struct PadConfig { pub padding: Vec<[usize; 2]>, } impl PadChip { pub fn param_vec_to_config(layer_params: Vec) -> PadConfig { assert!(layer_params.len() % 2 == 0); let padding = layer_params .chunks(2) .map(|chunk| [chunk[0] as usize, chunk[1] as usize]) .collect(); PadConfig { padding } } } impl Layer for PadChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { // FIXME: the pad from tflite is actually two, but mine is one // assert_eq!(tensors.len(), 1); let input = &tensors[0]; let zero = constants.get(&0).unwrap().clone(); let padding = PadChip::param_vec_to_config(layer_config.layer_params.clone()); let padded = pad(input, padding.padding, &zero); Ok(vec![padded]) } } impl GadgetConsumer for PadChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/permute.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::IxDyn; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct PermuteChip {} impl Layer for PermuteChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let params = &layer_config .layer_params .iter() .map(|x| *x as usize) .collect::>()[..]; assert!(inp.ndim() == params.len()); let out = inp.clone(); let out = out.permuted_axes(IxDyn(params)); Ok(vec![out]) } } impl GadgetConsumer for PermuteChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/reshape.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::Array; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct ReshapeChip {} impl Layer for ReshapeChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let shape = layer_config.out_shapes[0].clone(); println!("Reshape: {:?} -> {:?}", inp.shape(), shape); let flat = inp.iter().map(|x| x.clone()).collect(); let out = Array::from_shape_vec(shape, flat).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for ReshapeChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/resize_nn.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct ResizeNNChip {} // TODO: this does not work in general impl Layer for ResizeNNChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let output_shape = layer_config.out_shapes[0].clone(); assert_eq!(inp.ndim(), 4); assert_eq!(inp.shape()[0], 1); assert_eq!(inp.shape()[3], output_shape[3]); let mut flat = vec![]; // Do nearest neighbor interpolation over batch, h, w, c // The interpolation is over h and w for b in 0..inp.shape()[0] { for h in 0..output_shape[1] { let h_in = (h as f64 * (inp.shape()[1] as f64 / output_shape[1] as f64)) as usize; for w in 0..output_shape[2] { let w_in = (w as f64 * (inp.shape()[2] as f64 / output_shape[2] as f64)) as usize; for c in 0..inp.shape()[3] { flat.push(inp[[b, h_in, w_in, c]].clone()); } } } } let outp = Array::from_shape_vec(IxDyn(&output_shape), flat).unwrap(); Ok(vec![outp]) } } impl GadgetConsumer for ResizeNNChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/rotate.rs ================================================ // TODO: The implementation is not ideal. use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct RotateChip {} // Example: // input: // [1 2 3 4] // [5 6 7 8] // // params: [1] -- flip axis 1 only // output: // [4 3 2 1] // [8 7 6 5] impl Layer for RotateChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let params = &layer_config.layer_params; assert!(inp.shape().len() == 4); let mut flip = vec![false; 4]; for p in params { flip[*p as usize] = true; } let shape = inp.shape(); println!("Rotate: {:?} -> {:?}", inp.shape(), shape); let mut out = inp.clone(); for i in 0..shape[0] { for j in 0..shape[1] { for k in 0..shape[2] { for l in 0..shape[3] { let [ix, jx, kx, lx]: [usize; 4] = [i, j, k, l] .iter() .enumerate() .map(|(idx, x)| if flip[idx] { shape[idx] - 1 - *x } else { *x }) .collect::>() .try_into() .unwrap(); out[[ix, jx, kx, lx]] = inp[[i, j, k, l]].clone(); } } } } Ok(vec![out]) } } impl GadgetConsumer for RotateChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/slice.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::Slice; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct SliceChip {} impl Layer for SliceChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let params = &layer_config.layer_params; assert_eq!(params.len() % 2, 0); let num_axes = params.len() / 2; let starts = ¶ms[0..num_axes]; let sizes = ¶ms[num_axes..]; let inp = &tensors[0]; let outp = inp.slice_each_axis(|ax| { let start = starts[ax.axis.0] as usize; let size = sizes[ax.axis.0]; if size == -1 { Slice::from(start..) } else { Slice::from(start..(start + size as usize)) } }); Ok(vec![outp.to_owned()]) } } impl GadgetConsumer for SliceChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/split.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Axis, Slice}; use crate::{ gadgets::gadget::{GadgetConfig, GadgetType}, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct SplitChip {} impl Layer for SplitChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let axis = layer_config.layer_params[0] as usize; let num_splits = layer_config.layer_params[1] as usize; let inp = &tensors[1]; let mut out = vec![]; let split_len = inp.shape()[axis] / num_splits; for i in 0..num_splits { let slice = inp .slice_axis( Axis(axis), Slice::from((i * split_len)..((i + 1) * split_len)), ) .to_owned(); out.push(slice.to_owned()); } Ok(out) } } impl GadgetConsumer for SplitChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape/transpose.rs ================================================ use std::{collections::HashMap, rc::Rc}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::{ gadgets::gadget::GadgetConfig, layers::layer::{AssignedTensor, CellRc, GadgetConsumer}, }; use super::super::layer::{Layer, LayerConfig}; pub struct TransposeChip {} impl Layer for TransposeChip { fn forward( &self, _layouter: impl Layouter, tensors: &Vec>, _constants: &HashMap>, _gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { assert_eq!(layer_config.layer_params.len() % 2, 0); let ndim = layer_config.layer_params.len() / 2; let inp_shape = layer_config.layer_params[0..ndim] .to_vec() .iter() .map(|x| *x as usize) .collect::>(); let permutation = layer_config.layer_params[ndim..] .to_vec() .iter() .map(|x| *x as usize) .collect::>(); let inp = &tensors[0]; // Required because of memory layout issues let inp_flat = inp.iter().cloned().collect::>(); let inp = Array::from_shape_vec(IxDyn(&inp_shape), inp_flat).unwrap(); let inp = inp.permuted_axes(IxDyn(&permutation)); Ok(vec![inp]) } } impl GadgetConsumer for TransposeChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![] } } ================================================ FILE: src/layers/shape.rs ================================================ pub mod broadcast; pub mod concatenation; pub mod mask_neg_inf; pub mod pack; pub mod pad; pub mod permute; pub mod reshape; pub mod resize_nn; pub mod rotate; pub mod slice; pub mod split; pub mod transpose; ================================================ FILE: src/layers/softmax.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{ circuit::{AssignedCell, Layouter}, halo2curves::ff::PrimeField, plonk::Error, }; use ndarray::{s, Array, IxDyn}; use crate::gadgets::{ adder::AdderChip, gadget::{Gadget, GadgetConfig, GadgetType}, max::MaxChip, nonlinear::exp::ExpGadgetChip, sub_pairs::SubPairsChip, var_div_big3::VarDivRoundBig3Chip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct SoftmaxChip {} impl SoftmaxChip { pub fn softmax_flat( mut layouter: impl Layouter, constants: &HashMap>, inp_flat: Vec<&AssignedCell>, gadget_config: Rc, mask: &Vec, ) -> Result>, Error> { let exp_chip = ExpGadgetChip::::construct(gadget_config.clone()); let adder_chip = AdderChip::::construct(gadget_config.clone()); let sub_pairs_chip = SubPairsChip::::construct(gadget_config.clone()); let max_chip = MaxChip::::construct(gadget_config.clone()); let var_div_big_chip = VarDivRoundBig3Chip::::construct(gadget_config.clone()); let zero = constants.get(&0).unwrap().as_ref(); let sf = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); // Mask the input for max computation and subtraction let inp_take = inp_flat .iter() .enumerate() .filter(|(i, _)| mask[*i] == 0) // Awkwardly, 1 = take negative infinity .map(|(_, x)| *x) .collect::>(); // Compute the max let max = max_chip .forward( layouter.namespace(|| format!("max")), &vec![inp_take.clone()], &vec![zero], ) .unwrap(); let max = &max[0]; // Subtract the max let max_flat = vec![max; inp_take.len()]; let sub = sub_pairs_chip.forward( layouter.namespace(|| format!("sub")), &vec![inp_take, max_flat], &vec![zero], )?; let sub = sub.iter().collect::>(); // Compute the exp let exp_slice = exp_chip.forward( layouter.namespace(|| format!("exp")), &vec![sub], &vec![zero], )?; // Compute the sum let sum = adder_chip.forward( layouter.namespace(|| format!("sum")), &vec![exp_slice.iter().collect()], &vec![zero], )?; let sum = sum[0].clone(); let sum_div_sf = var_div_big_chip.forward( layouter.namespace(|| format!("sum div sf")), &vec![vec![&sum]], &vec![zero, sf], )?; let sum_div_sf = sum_div_sf[0].clone(); let dived = var_div_big_chip.forward( layouter.namespace(|| format!("div")), &vec![exp_slice.iter().collect()], &vec![zero, &sum_div_sf], )?; // Take either zero (softmax(-inf)) or the result let mut div_idx = 0; let dived = mask .iter() .map(|x| { if *x == 1 { zero.clone() } else { let tmp = dived[div_idx].clone(); div_idx = div_idx + 1; tmp } }) .collect(); Ok(dived) } } impl Layer for SoftmaxChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; assert!(inp.ndim() == 2 || inp.ndim() == 3 || inp.ndim() == 4); if inp.ndim() == 4 { assert_eq!(inp.shape()[0], 1); } let inp_shape = inp.shape().iter().map(|x| *x).collect::>(); let mask = if layer_config.layer_params.len() == 0 { Array::from_shape_fn(IxDyn(&inp_shape), |_| 0) } else { let mask_shape_len = layer_config.layer_params[0] as usize; let mask_shape = layer_config.layer_params[1..(1 + mask_shape_len)] .iter() .map(|x| *x as usize) .collect::>(); let mask = layer_config.layer_params[(1 + mask_shape_len)..].to_vec(); let mask = Array::from_shape_vec(IxDyn(&mask_shape), mask).unwrap(); let mask = mask.broadcast(IxDyn(&inp_shape)).unwrap().to_owned(); mask }; let shape = if inp.ndim() == 2 || inp.ndim() == 3 { inp.shape().iter().map(|x| *x).collect::>() } else { vec![inp.shape()[1], inp.shape()[2], inp.shape()[3]] }; let inp = inp.to_owned().into_shape(shape.clone()).unwrap(); let mask = mask.into_shape(shape.clone()).unwrap(); let mut outp = vec![]; if inp.ndim() == 2 { for i in 0..shape[0] { let inp_slice = inp.slice(s![i, ..]); let inp_flat = inp_slice.iter().map(|x| x.as_ref()).collect::>(); let mask_slice = mask.slice(s![i, ..]); let mask_flat = mask_slice.iter().map(|x| *x as i64).collect::>(); let dived = Self::softmax_flat( layouter.namespace(|| format!("softmax {}", i)), constants, inp_flat, gadget_config.clone(), &mask_flat, ) .unwrap(); outp.extend(dived); } } else if inp.ndim() == 3 { for i in 0..shape[0] { for j in 0..shape[1] { let inp_slice = inp.slice(s![i, j, ..]); let inp_flat = inp_slice.iter().map(|x| x.as_ref()).collect::>(); let mask_slice = mask.slice(s![i, j, ..]); let mask_flat = mask_slice.iter().map(|x| *x as i64).collect::>(); let dived = Self::softmax_flat( layouter.namespace(|| format!("softmax {} {}", i, j)), constants, inp_flat, gadget_config.clone(), &mask_flat, ) .unwrap(); outp.extend(dived); } } } else { panic!("Not implemented"); } let outp = outp.into_iter().map(|x| Rc::new(x)).collect::>(); let outp = Array::from_shape_vec(IxDyn(inp.shape()), outp).unwrap(); Ok(vec![outp]) } } impl GadgetConsumer for SoftmaxChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::Exp, GadgetType::Adder, GadgetType::VarDivRoundBig3, GadgetType::Max, GadgetType::SubPairs, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/sqrt.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::sqrt::SqrtGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct SqrtChip {} impl Layer for SqrtChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let mut inp_vec = vec![]; let mask = &layer_config.mask; let mut mask_map = HashMap::new(); for i in 0..mask.len() / 2 { mask_map.insert(mask[2 * i], mask[2 * i + 1]); } let min_val = gadget_config.min_val; let min_val = constants.get(&min_val).unwrap().as_ref(); let max_val = gadget_config.max_val; let max_val = constants.get(&max_val).unwrap().as_ref(); for (i, val) in inp.iter().enumerate() { let i = i as i64; if mask_map.contains_key(&i) { let mask_val = *mask_map.get(&i).unwrap(); if mask_val == 1 { inp_vec.push(max_val); } else if mask_val == -1 { inp_vec.push(min_val); } else { panic!(); } } else { inp_vec.push(val.as_ref()); } } let zero = constants.get(&0).unwrap().as_ref(); let sqrt_chip = SqrtGadgetChip::::construct(gadget_config.clone()); let vec_inps = vec![inp_vec]; let constants = vec![zero, min_val, max_val]; let out = sqrt_chip.forward(layouter.namespace(|| "sqrt chip"), &vec_inps, &constants)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for SqrtChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Sqrt, GadgetType::InputLookup] } } ================================================ FILE: src/layers/square.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, square::SquareGadgetChip, var_div::VarDivRoundChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct SquareChip {} impl Layer for SquareChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { assert_eq!(tensors.len(), 1); let inp = &tensors[0]; let zero = constants.get(&0).unwrap().as_ref(); let square_chip = SquareGadgetChip::::construct(gadget_config.clone()); let inp_vec = inp.iter().map(|x| x.as_ref()).collect::>(); let vec_inputs = vec![inp_vec]; let single_inps = vec![zero]; let out = square_chip.forward( layouter.namespace(|| "square chip"), &vec_inputs, &single_inps, )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let div = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); let single_inps = vec![zero, div]; let out = out.iter().collect::>(); let vec_inputs = vec![out]; let out = var_div_chip.forward( layouter.namespace(|| "var div chip"), &vec_inputs, &single_inps, )?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for SquareChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::Square, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/squared_diff.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::{ gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, squared_diff::SquaredDiffGadgetChip, var_div::VarDivRoundChip, }, utils::helpers::broadcast, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct SquaredDiffChip {} impl Layer for SquaredDiffChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { assert_eq!(tensors.len(), 2); let inp1 = &tensors[0]; let inp2 = &tensors[1]; // Broadcoasting allowed... can't check shapes easily let (inp1, inp2) = broadcast(inp1, inp2); let zero = constants.get(&0).unwrap().as_ref(); let sq_diff_chip = SquaredDiffGadgetChip::::construct(gadget_config.clone()); let inp1_vec = inp1.iter().map(|x| x.as_ref()).collect::>(); let inp2_vec = inp2.iter().map(|x| x.as_ref()).collect::>(); let vec_inputs = vec![inp1_vec, inp2_vec]; let tmp_constants = vec![zero]; let out = sq_diff_chip.forward( layouter.namespace(|| "sq diff chip"), &vec_inputs, &tmp_constants, )?; let var_div_chip = VarDivRoundChip::::construct(gadget_config.clone()); let div = constants .get(&(gadget_config.scale_factor as i64)) .unwrap() .as_ref(); let single_inputs = vec![zero, div]; let out = out.iter().map(|x| x).collect::>(); let out = var_div_chip.forward( layouter.namespace(|| "sq diff div"), &vec![out], &single_inputs, )?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp1.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for SquaredDiffChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![ GadgetType::SquaredDiff, GadgetType::VarDivRound, GadgetType::InputLookup, ] } } ================================================ FILE: src/layers/tanh.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, nonlinear::tanh::TanhGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct TanhChip {} impl Layer for TanhChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let inp = &tensors[0]; let inp_vec = inp.iter().map(|x| x.as_ref()).collect::>(); let zero = constants.get(&0).unwrap().as_ref(); let tanh_chip = TanhGadgetChip::::construct(gadget_config.clone()); let vec_inps = vec![inp_vec]; let constants = vec![zero]; let out = tanh_chip.forward(layouter.namespace(|| "tanh chip"), &vec_inps, &constants)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(inp.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for TanhChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Tanh, GadgetType::InputLookup] } } ================================================ FILE: src/layers/update.rs ================================================ use std::{collections::HashMap, rc::Rc, vec}; use halo2_proofs::{circuit::Layouter, halo2curves::ff::PrimeField, plonk::Error}; use ndarray::{Array, IxDyn}; use crate::gadgets::{ gadget::{Gadget, GadgetConfig, GadgetType}, update::UpdateGadgetChip, }; use super::layer::{AssignedTensor, CellRc, GadgetConsumer, Layer, LayerConfig}; #[derive(Clone, Debug)] pub struct UpdateChip {} impl Layer for UpdateChip { fn forward( &self, mut layouter: impl Layouter, tensors: &Vec>, constants: &HashMap>, gadget_config: Rc, _layer_config: &LayerConfig, ) -> Result>, Error> { let w = &tensors[0]; let dw = &tensors[1]; let zero = constants.get(&0).unwrap().as_ref(); let update_chip = UpdateGadgetChip::::construct((*gadget_config).clone()); let flattened_w = w.into_iter().map(|x| (**x).clone()).collect::>(); let flattened_dw = dw.into_iter().map(|x| (**x).clone()).collect::>(); let flattened_w_ref = flattened_w.iter().collect::>(); let flattened_dw_ref = flattened_dw.iter().collect::>(); let vec_inps = vec![flattened_w_ref, flattened_dw_ref]; let constants = vec![zero]; let out = update_chip.forward(layouter.namespace(|| "update chip"), &vec_inps, &constants)?; let out = out.into_iter().map(|x| Rc::new(x)).collect::>(); let out = Array::from_shape_vec(IxDyn(w.shape()), out).unwrap(); Ok(vec![out]) } } impl GadgetConsumer for UpdateChip { fn used_gadgets(&self, _layer_params: Vec) -> Vec { vec![GadgetType::Update] } } ================================================ FILE: src/layers.rs ================================================ // Generics pub mod averager; pub mod arithmetic; pub mod shape; // Concrete implementations pub mod avg_pool_2d; pub mod batch_mat_mul; pub mod conv2d; pub mod div_fixed; pub mod fully_connected; pub mod logistic; pub mod max_pool_2d; pub mod mean; pub mod noop; pub mod pow; pub mod rsqrt; pub mod softmax; pub mod sqrt; pub mod square; pub mod squared_diff; pub mod tanh; pub mod update; // Special: dag pub mod dag; // Special: layer pub mod layer; ================================================ FILE: src/lib.rs ================================================ #![feature(int_roundings)] pub mod commitments; pub mod gadgets; pub mod layers; pub mod model; pub mod utils; ================================================ FILE: src/model.rs ================================================ use std::{ collections::{BTreeMap, BTreeSet, HashMap}, marker::PhantomData, rc::Rc, sync::{Arc, Mutex}, }; use halo2_proofs::{ circuit::{Layouter, SimpleFloorPlanner, Value}, halo2curves::ff::{FromUniformBytes, PrimeField}, plonk::{Advice, Circuit, Column, ConstraintSystem, Error, Instance}, }; use lazy_static::lazy_static; use ndarray::{Array, IxDyn}; use num_bigint::BigUint; use crate::{ commitments::{ commit::Commit, packer::PackerChip, poseidon_commit::{PoseidonCommitChip, L, RATE, WIDTH}, }, gadgets::{ add_pairs::AddPairsChip, adder::AdderChip, bias_div_round_relu6::BiasDivRoundRelu6Chip, dot_prod::DotProductChip, gadget::{Gadget, GadgetConfig, GadgetType}, input_lookup::InputLookupChip, max::MaxChip, mul_pairs::MulPairsChip, nonlinear::{exp::ExpGadgetChip, pow::PowGadgetChip, relu::ReluChip, tanh::TanhGadgetChip}, nonlinear::{logistic::LogisticGadgetChip, rsqrt::RsqrtGadgetChip, sqrt::SqrtGadgetChip}, sqrt_big::SqrtBigChip, square::SquareGadgetChip, squared_diff::SquaredDiffGadgetChip, sub_pairs::SubPairsChip, update::UpdateGadgetChip, var_div::VarDivRoundChip, var_div_big::VarDivRoundBigChip, var_div_big3::VarDivRoundBig3Chip, }, layers::{ arithmetic::{add::AddChip, div_var::DivVarChip, mul::MulChip, sub::SubChip}, avg_pool_2d::AvgPool2DChip, batch_mat_mul::BatchMatMulChip, conv2d::Conv2DChip, dag::{DAGLayerChip, DAGLayerConfig}, fully_connected::{FullyConnectedChip, FullyConnectedConfig}, layer::{AssignedTensor, CellRc, GadgetConsumer, LayerConfig, LayerType}, logistic::LogisticChip, max_pool_2d::MaxPool2DChip, mean::MeanChip, noop::NoopChip, pow::PowChip, rsqrt::RsqrtChip, shape::{ broadcast::BroadcastChip, concatenation::ConcatenationChip, mask_neg_inf::MaskNegInfChip, pack::PackChip, pad::PadChip, permute::PermuteChip, reshape::ReshapeChip, resize_nn::ResizeNNChip, rotate::RotateChip, slice::SliceChip, split::SplitChip, transpose::TransposeChip, }, softmax::SoftmaxChip, sqrt::SqrtChip, square::SquareChip, squared_diff::SquaredDiffChip, tanh::TanhChip, update::UpdateChip, }, utils::{ helpers::{convert_to_bigint, RAND_START_IDX}, loader::{load_model_msgpack, ModelMsgpack}, }, }; lazy_static! { pub static ref GADGET_CONFIG: Mutex = Mutex::new(GadgetConfig::default()); pub static ref PUBLIC_VALS: Mutex> = Mutex::new(vec![]); } #[derive(Clone, Debug, Default)] pub struct ModelCircuit { pub used_gadgets: Arc>, pub dag_config: DAGLayerConfig, pub tensors: BTreeMap>, pub commit_before: Vec>, pub commit_after: Vec>, pub k: usize, pub bits_per_elem: usize, pub inp_idxes: Vec, pub num_random: i64, } #[derive(Clone, Debug)] pub struct ModelConfig> { pub gadget_config: Rc, pub public_col: Column, pub hasher: Option>, pub _marker: PhantomData, } impl> ModelCircuit { pub fn assign_tensors_map( &self, mut layouter: impl Layouter, columns: &Vec>, tensors: &BTreeMap>, ) -> Result>, Error> { let tensors = layouter.assign_region( || "asssignment", |mut region| { let mut cell_idx = 0; let mut assigned_tensors = BTreeMap::new(); for (tensor_idx, tensor) in tensors.iter() { let mut flat = vec![]; for val in tensor.iter() { let row_idx = cell_idx / columns.len(); let col_idx = cell_idx % columns.len(); let cell = region .assign_advice( || "assignment", columns[col_idx], row_idx, || Value::known(*val), ) .unwrap(); flat.push(Rc::new(cell)); cell_idx += 1; } let tensor = Array::from_shape_vec(tensor.shape(), flat).unwrap(); assigned_tensors.insert(*tensor_idx, tensor); } Ok(assigned_tensors) }, )?; Ok(tensors) } pub fn tensor_map_to_vec( &self, tensor_map: &BTreeMap, IxDyn>>, ) -> Result>, Error> { let smallest_tensor = tensor_map .iter() .min_by_key(|(_, tensor)| tensor.len()) .unwrap() .1; let max_tensor_key = tensor_map .iter() .max_by_key(|(key, _)| *key) .unwrap() .0 .clone(); let mut tensors = vec![]; for i in 0..max_tensor_key + 1 { let tensor = tensor_map.get(&i).unwrap_or(smallest_tensor); tensors.push(tensor.clone()); } Ok(tensors) } pub fn assign_tensors_vec( &self, mut layouter: impl Layouter, columns: &Vec>, tensors: &BTreeMap>, ) -> Result>, Error> { let tensor_map = self .assign_tensors_map( layouter.namespace(|| "assign_tensors_map"), columns, tensors, ) .unwrap(); self.tensor_map_to_vec(&tensor_map) } pub fn assign_constants( &self, mut layouter: impl Layouter, gadget_config: Rc, ) -> Result>, Error> { let sf = gadget_config.scale_factor; let min_val = gadget_config.min_val; let max_val = gadget_config.max_val; let constants = layouter.assign_region( || "constants", |mut region| { let mut constants: HashMap> = HashMap::new(); let vals = vec![0 as i64, 1, sf as i64, min_val, max_val]; let shift_val_i64 = -min_val * 2; // FIXME let shift_val_f = F::from(shift_val_i64 as u64); for (i, val) in vals.iter().enumerate() { let cell = region.assign_fixed( || format!("constant_{}", i), gadget_config.fixed_columns[0], i, || Value::known(F::from((val + shift_val_i64) as u64) - shift_val_f), )?; constants.insert(*val, Rc::new(cell)); } // TODO: I've made some very bad life decisions // TOOD: this needs to be a random oracle let r_base = F::from(0x123456789abcdef); let mut r = r_base.clone(); for i in 0..self.num_random { let rand = region.assign_fixed( || format!("rand_{}", i), gadget_config.fixed_columns[0], constants.len(), || Value::known(r), )?; r = r * r_base; constants.insert(RAND_START_IDX + (i as i64), Rc::new(rand)); } Ok(constants) }, )?; Ok(constants) } // TODO: for some horrifying reason, assigning to fixed columns causes everything to blow up // Currently get around this by assigning to advice columns // This is secure because of the equality checks but EXTREMELY STUPID pub fn assign_constants2( &self, mut layouter: impl Layouter, gadget_config: Rc, fixed_constants: &HashMap>, ) -> Result>, Error> { let sf = gadget_config.scale_factor; let min_val = gadget_config.min_val; let max_val = gadget_config.max_val; let constants = layouter.assign_region( || "constants", |mut region| { let mut constants: HashMap> = HashMap::new(); let vals = vec![0 as i64, 1, sf as i64, min_val, max_val]; let shift_val_i64 = -min_val * 2; // FIXME let shift_val_f = F::from(shift_val_i64 as u64); for (i, val) in vals.iter().enumerate() { let assignment_idx = i as usize; let row_idx = assignment_idx / gadget_config.columns.len(); let col_idx = assignment_idx % gadget_config.columns.len(); let cell = region.assign_advice( || format!("constant_{}", i), gadget_config.columns[col_idx], row_idx, || Value::known(F::from((val + shift_val_i64) as u64) - shift_val_f), )?; constants.insert(*val, Rc::new(cell)); } // TODO: I've made some very bad life decisions // TOOD: this needs to be a random oracle let r_base = F::from(0x123456789abcdef); let mut r = r_base.clone(); for i in 0..self.num_random { let assignment_idx = constants.len(); let row_idx = assignment_idx / gadget_config.columns.len(); let col_idx = assignment_idx % gadget_config.columns.len(); let rand = region.assign_advice( || format!("rand_{}", i), gadget_config.columns[col_idx], row_idx, || Value::known(r), )?; r = r * r_base; constants.insert(RAND_START_IDX + (i as i64), Rc::new(rand)); } for (k, v) in fixed_constants.iter() { let v2 = constants.get(k).unwrap(); region.constrain_equal(v.cell(), v2.cell()).unwrap(); } Ok(constants) }, )?; Ok(constants) } pub fn generate_from_file(config_file: &str, inp_file: &str) -> ModelCircuit { let config = load_model_msgpack(config_file, inp_file); Self::generate_from_msgpack(config, true) } pub fn generate_from_msgpack(config: ModelMsgpack, panic_empty_tensor: bool) -> ModelCircuit { let to_field = |x: i64| { let bias = 1 << 31; let x_pos = x + bias; F::from(x_pos as u64) - F::from(bias as u64) }; let match_layer = |x: &str| match x { "AveragePool2D" => LayerType::AvgPool2D, "Add" => LayerType::Add, "BatchMatMul" => LayerType::BatchMatMul, "Broadcast" => LayerType::Broadcast, "Concatenation" => LayerType::Concatenation, "Conv2D" => LayerType::Conv2D, "Div" => LayerType::DivFixed, // TODO: rename to DivFixed "DivVar" => LayerType::DivVar, "FullyConnected" => LayerType::FullyConnected, "Logistic" => LayerType::Logistic, "MaskNegInf" => LayerType::MaskNegInf, "MaxPool2D" => LayerType::MaxPool2D, "Mean" => LayerType::Mean, "Mul" => LayerType::Mul, "Noop" => LayerType::Noop, "Pack" => LayerType::Pack, "Pad" => LayerType::Pad, "Pow" => LayerType::Pow, "Permute" => LayerType::Permute, "Reshape" => LayerType::Reshape, "ResizeNearestNeighbor" => LayerType::ResizeNN, "Rotate" => LayerType::Rotate, "Rsqrt" => LayerType::Rsqrt, "Slice" => LayerType::Slice, "Softmax" => LayerType::Softmax, "Split" => LayerType::Split, "Sqrt" => LayerType::Sqrt, "Square" => LayerType::Square, "SquaredDifference" => LayerType::SquaredDifference, "Sub" => LayerType::Sub, "Tanh" => LayerType::Tanh, "Transpose" => LayerType::Transpose, "Update" => LayerType::Update, _ => panic!("unknown op: {}", x), }; let mut tensors = BTreeMap::new(); for flat in config.tensors { let value_flat = flat.data.iter().map(|x| to_field(*x)).collect::>(); let shape = flat.shape.iter().map(|x| *x as usize).collect::>(); let num_el: usize = shape.iter().product(); if panic_empty_tensor && num_el != value_flat.len() { panic!("tensor shape and data length mismatch"); } if num_el == value_flat.len() { let tensor = Array::from_shape_vec(IxDyn(&shape), value_flat).unwrap(); tensors.insert(flat.idx, tensor); } else { // Do nothing here since we're loading the config }; } let i64_to_usize = |x: &Vec| x.iter().map(|x| *x as usize).collect::>(); let mut used_gadgets = BTreeSet::new(); let dag_config = { let ops = config .layers .iter() .map(|layer| { let layer_type = match_layer(&layer.layer_type); let layer_gadgets = match layer_type { LayerType::Add => Box::new(AddChip {}) as Box, LayerType::AvgPool2D => Box::new(AvgPool2DChip {}) as Box, LayerType::BatchMatMul => Box::new(BatchMatMulChip {}) as Box, LayerType::Broadcast => Box::new(BroadcastChip {}) as Box, LayerType::Concatenation => Box::new(ConcatenationChip {}) as Box, LayerType::DivFixed => Box::new(ConcatenationChip {}) as Box, LayerType::DivVar => Box::new(DivVarChip {}) as Box, LayerType::Conv2D => Box::new(Conv2DChip { config: LayerConfig::default(), _marker: PhantomData::, }) as Box, LayerType::FullyConnected => Box::new(FullyConnectedChip { config: FullyConnectedConfig { normalize: true }, _marker: PhantomData::, }) as Box, LayerType::Logistic => Box::new(LogisticChip {}) as Box, LayerType::MaskNegInf => Box::new(MaskNegInfChip {}) as Box, LayerType::MaxPool2D => Box::new(MaxPool2DChip { marker: PhantomData::, }) as Box, LayerType::Mean => Box::new(MeanChip {}) as Box, LayerType::Mul => Box::new(MulChip {}) as Box, LayerType::Noop => Box::new(NoopChip {}) as Box, LayerType::Pack => Box::new(PackChip {}) as Box, LayerType::Pad => Box::new(PadChip {}) as Box, LayerType::Pow => Box::new(PowChip {}) as Box, LayerType::Permute => Box::new(PermuteChip {}) as Box, LayerType::Reshape => Box::new(ReshapeChip {}) as Box, LayerType::ResizeNN => Box::new(ResizeNNChip {}) as Box, LayerType::Rotate => Box::new(RotateChip {}) as Box, LayerType::Rsqrt => Box::new(RsqrtChip {}) as Box, LayerType::Slice => Box::new(SliceChip {}) as Box, LayerType::Softmax => Box::new(SoftmaxChip {}) as Box, LayerType::Split => Box::new(SplitChip {}) as Box, LayerType::Sqrt => Box::new(SqrtChip {}) as Box, LayerType::Square => Box::new(SquareChip {}) as Box, LayerType::SquaredDifference => Box::new(SquaredDiffChip {}) as Box, LayerType::Sub => Box::new(SubChip {}) as Box, LayerType::Tanh => Box::new(TanhChip {}) as Box, LayerType::Transpose => Box::new(TransposeChip {}) as Box, LayerType::Update => Box::new(UpdateChip {}) as Box, } .used_gadgets(layer.params.clone()); for gadget in layer_gadgets { used_gadgets.insert(gadget); } LayerConfig { layer_type, layer_params: layer.params.clone(), inp_shapes: layer.inp_shapes.iter().map(|x| i64_to_usize(x)).collect(), out_shapes: layer.out_shapes.iter().map(|x| i64_to_usize(x)).collect(), mask: layer.mask.clone(), } }) .collect::>(); let inp_idxes = config .layers .iter() .map(|layer| i64_to_usize(&layer.inp_idxes)) .collect::>(); let out_idxes = config .layers .iter() .map(|layer| i64_to_usize(&layer.out_idxes)) .collect::>(); let final_out_idxes = config .out_idxes .iter() .map(|x| *x as usize) .collect::>(); DAGLayerConfig { inp_idxes, out_idxes, ops, final_out_idxes, } }; // The input lookup is always used used_gadgets.insert(GadgetType::InputLookup); let used_gadgets = Arc::new(used_gadgets); let gadget = &GADGET_CONFIG; let cloned_gadget = gadget.lock().unwrap().clone(); *gadget.lock().unwrap() = GadgetConfig { scale_factor: config.global_sf as u64, shift_min_val: -(config.global_sf * config.global_sf * (1 << 17)), div_outp_min_val: -(1 << (config.k - 1)), min_val: -(1 << (config.k - 1)), max_val: (1 << (config.k - 1)) - 10, k: config.k as usize, num_rows: (1 << config.k) - 10 + 1, num_cols: config.num_cols as usize, used_gadgets: used_gadgets.clone(), commit_before: config.commit_before.clone().unwrap_or(vec![]), commit_after: config.commit_after.clone().unwrap_or(vec![]), use_selectors: config.use_selectors.unwrap_or(true), num_bits_per_elem: config.bits_per_elem.unwrap_or(config.k), ..cloned_gadget }; ModelCircuit { tensors, dag_config, used_gadgets, k: config.k as usize, bits_per_elem: config.bits_per_elem.unwrap_or(config.k) as usize, inp_idxes: config.inp_idxes.clone(), commit_after: config.commit_after.unwrap_or(vec![]), commit_before: config.commit_before.unwrap_or(vec![]), num_random: config.num_random.unwrap_or(0), } } pub fn assign_and_commit( &self, mut layouter: impl Layouter, constants: &HashMap>, config: &ModelConfig, tensors: &BTreeMap>, ) -> (BTreeMap>, CellRc) { let num_bits = self.bits_per_elem; let packer_config = PackerChip::::construct(num_bits, config.gadget_config.as_ref()); let packer_chip = PackerChip:: { config: packer_config, }; let (tensor_map, packed) = packer_chip .assign_and_pack( layouter.namespace(|| "packer"), config.gadget_config.clone(), constants, tensors, ) .unwrap(); let zero = constants.get(&0).unwrap().clone(); let commit_chip = config.hasher.clone().unwrap(); let commitments = commit_chip .commit( layouter.namespace(|| "commit"), config.gadget_config.clone(), constants, &packed, zero.clone(), ) .unwrap(); assert_eq!(commitments.len(), 1); (tensor_map, commitments[0].clone()) } pub fn copy_and_commit( &self, mut layouter: impl Layouter, constants: &HashMap>, config: &ModelConfig, tensors: &BTreeMap>, ) -> CellRc { let num_bits = self.bits_per_elem; let packer_config = PackerChip::::construct(num_bits, config.gadget_config.as_ref()); let packer_chip = PackerChip:: { config: packer_config, }; let packed = packer_chip .copy_and_pack( layouter.namespace(|| "packer"), config.gadget_config.clone(), constants, tensors, ) .unwrap(); let zero = constants.get(&0).unwrap().clone(); let commit_chip = config.hasher.clone().unwrap(); let commitments = commit_chip .commit( layouter.namespace(|| "commit"), config.gadget_config.clone(), constants, &packed, zero.clone(), ) .unwrap(); assert_eq!(commitments.len(), 1); commitments[0].clone() } } impl> Circuit for ModelCircuit { type Config = ModelConfig; type FloorPlanner = SimpleFloorPlanner; type Params = (); fn without_witnesses(&self) -> Self { todo!() } fn configure(meta: &mut ConstraintSystem) -> Self::Config { let mut gadget_config = crate::model::GADGET_CONFIG.lock().unwrap().clone(); let columns = (0..gadget_config.num_cols) .map(|_| meta.advice_column()) .collect::>(); for col in columns.iter() { meta.enable_equality(*col); } gadget_config.columns = columns; let public_col = meta.instance_column(); meta.enable_equality(public_col); gadget_config.fixed_columns = vec![meta.fixed_column()]; meta.enable_equality(gadget_config.fixed_columns[0]); // The input lookup is always loaded gadget_config = InputLookupChip::::configure(meta, gadget_config); let used_gadgets = gadget_config.used_gadgets.clone(); for gadget_type in used_gadgets.iter() { gadget_config = match gadget_type { GadgetType::AddPairs => AddPairsChip::::configure(meta, gadget_config), GadgetType::Adder => AdderChip::::configure(meta, gadget_config), GadgetType::BiasDivRoundRelu6 => BiasDivRoundRelu6Chip::::configure(meta, gadget_config), GadgetType::BiasDivFloorRelu6 => panic!(), GadgetType::DotProduct => DotProductChip::::configure(meta, gadget_config), GadgetType::Exp => ExpGadgetChip::::configure(meta, gadget_config), GadgetType::Logistic => LogisticGadgetChip::::configure(meta, gadget_config), GadgetType::Max => MaxChip::::configure(meta, gadget_config), GadgetType::MulPairs => MulPairsChip::::configure(meta, gadget_config), GadgetType::Pow => PowGadgetChip::::configure(meta, gadget_config), GadgetType::Relu => ReluChip::::configure(meta, gadget_config), GadgetType::Rsqrt => RsqrtGadgetChip::::configure(meta, gadget_config), GadgetType::Sqrt => SqrtGadgetChip::::configure(meta, gadget_config), GadgetType::SqrtBig => SqrtBigChip::::configure(meta, gadget_config), GadgetType::Square => SquareGadgetChip::::configure(meta, gadget_config), GadgetType::SquaredDiff => SquaredDiffGadgetChip::::configure(meta, gadget_config), GadgetType::SubPairs => SubPairsChip::::configure(meta, gadget_config), GadgetType::Tanh => TanhGadgetChip::::configure(meta, gadget_config), GadgetType::VarDivRound => VarDivRoundChip::::configure(meta, gadget_config), GadgetType::VarDivRoundBig => VarDivRoundBigChip::::configure(meta, gadget_config), GadgetType::VarDivRoundBig3 => VarDivRoundBig3Chip::::configure(meta, gadget_config), GadgetType::InputLookup => gadget_config, // This is always loaded GadgetType::Update => UpdateGadgetChip::::configure(meta, gadget_config), GadgetType::Packer => panic!(), }; } let hasher = if gadget_config.commit_before.len() + gadget_config.commit_after.len() > 0 { let packer_config = PackerChip::::construct(gadget_config.num_bits_per_elem as usize, &gadget_config); gadget_config = PackerChip::::configure(meta, packer_config, gadget_config); // TODO let input = gadget_config.columns[0..L].try_into().unwrap(); let state = gadget_config.columns[L..L + WIDTH].try_into().unwrap(); let partial_sbox = gadget_config.columns[L + WIDTH].into(); Some(PoseidonCommitChip::::configure( meta, input, state, partial_sbox, )) } else { None }; ModelConfig { gadget_config: gadget_config.into(), public_col, hasher, _marker: PhantomData, } } fn synthesize(&self, config: Self::Config, mut layouter: impl Layouter) -> Result<(), Error> { // Assign tables let gadget_rc: Rc = config.gadget_config.clone().into(); for gadget in self.used_gadgets.iter() { match gadget { GadgetType::AddPairs => { let chip = AddPairsChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "add pairs lookup"))?; } GadgetType::Adder => { let chip = AdderChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "adder lookup"))?; } GadgetType::BiasDivRoundRelu6 => { let chip = BiasDivRoundRelu6Chip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "bias div round relu6 lookup"))?; } GadgetType::DotProduct => { let chip = DotProductChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "dot product lookup"))?; } GadgetType::VarDivRound => { let chip = VarDivRoundChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "var div lookup"))?; } GadgetType::Pow => { let chip = PowGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "pow lookup"))?; } GadgetType::Relu => { let chip = ReluChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "relu lookup"))?; } GadgetType::Rsqrt => { let chip = RsqrtGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "rsqrt lookup"))?; } GadgetType::Sqrt => { let chip = SqrtGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "sqrt lookup"))?; } GadgetType::Tanh => { let chip = TanhGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "tanh lookup"))?; } GadgetType::Exp => { let chip = ExpGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "exp lookup"))?; } GadgetType::Logistic => { let chip = LogisticGadgetChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "logistic lookup"))?; } GadgetType::InputLookup => { let chip = InputLookupChip::::construct(gadget_rc.clone()); chip.load_lookups(layouter.namespace(|| "input lookup"))?; } GadgetType::VarDivRoundBig => {} GadgetType::VarDivRoundBig3 => {} GadgetType::Max => {} GadgetType::MulPairs => {} GadgetType::SqrtBig => {} GadgetType::Square => {} GadgetType::SquaredDiff => {} GadgetType::SubPairs => {} GadgetType::Update => {} _ => panic!("unsupported gadget {:?}", gadget), } } // Assign weights and constants let constants_base = self .assign_constants( layouter.namespace(|| "constants"), config.gadget_config.clone(), ) .unwrap(); // Some halo2 cancer let constants = self .assign_constants2( layouter.namespace(|| "constants 2"), config.gadget_config.clone(), &constants_base, ) .unwrap(); let mut commitments = vec![]; let tensors = if self.commit_before.len() > 0 { // Commit to the tensors before the DAG let mut tensor_map = BTreeMap::new(); let mut ignore_idxes: Vec = vec![]; for commit_idxes in self.commit_before.iter() { let to_commit = BTreeMap::from_iter( commit_idxes .iter() .map(|idx| (*idx, self.tensors.get(idx).unwrap().clone())), ); let (mut committed_tensors, commitment) = self.assign_and_commit( layouter.namespace(|| "commit"), &constants, &config, &to_commit, ); commitments.push(commitment); tensor_map.append(&mut committed_tensors); ignore_idxes.extend(commit_idxes.iter()); } // Assign the remainder of the tensors let mut assign_map = BTreeMap::new(); for (idx, tensor) in self.tensors.iter() { if ignore_idxes.contains(idx) { continue; } assign_map.insert(*idx, tensor.clone()); } let mut remainder_tensor_map = self .assign_tensors_map( layouter.namespace(|| "assignment"), &config.gadget_config.columns, &assign_map, ) .unwrap(); // Merge the two maps tensor_map.append(&mut remainder_tensor_map); // Return the tensors self.tensor_map_to_vec(&tensor_map).unwrap() } else { self .assign_tensors_vec( layouter.namespace(|| "assignment"), &config.gadget_config.columns, &self.tensors, ) .unwrap() }; // Perform the dag let dag_chip = DAGLayerChip::::construct(self.dag_config.clone()); let (final_tensor_map, result) = dag_chip.forward( layouter.namespace(|| "dag"), &tensors, &constants, config.gadget_config.clone(), &LayerConfig::default(), )?; if self.commit_after.len() > 0 { for commit_idxes in self.commit_after.iter() { let to_commit = BTreeMap::from_iter(commit_idxes.iter().map(|idx| { ( *idx, final_tensor_map.get(&(*idx as usize)).unwrap().clone(), ) })); let commitment = self.copy_and_commit( layouter.namespace(|| "commit"), &constants, &config, &to_commit, ); commitments.push(commitment); } } let mut pub_layouter = layouter.namespace(|| "public"); let mut total_idx = 0; let mut new_public_vals = vec![]; for cell in commitments.iter() { pub_layouter .constrain_instance(cell.as_ref().cell(), config.public_col, total_idx) .unwrap(); let val = convert_to_bigint(cell.value().map(|x| x.to_owned())); new_public_vals.push(val); total_idx += 1; } for tensor in result { for cell in tensor.iter() { pub_layouter .constrain_instance(cell.as_ref().cell(), config.public_col, total_idx) .unwrap(); let val = convert_to_bigint(cell.value().map(|x| x.to_owned())); new_public_vals.push(val); total_idx += 1; } } *PUBLIC_VALS.lock().unwrap() = new_public_vals; Ok(()) } } ================================================ FILE: src/utils/helpers.rs ================================================ use halo2_proofs::{ circuit::{AssignedCell, Value}, halo2curves::ff::PrimeField, }; use ndarray::{Array, IxDyn}; use num_bigint::BigUint; use crate::{gadgets::gadget::convert_to_u128, model::PUBLIC_VALS}; // TODO: this is very bad pub const RAND_START_IDX: i64 = i64::MIN; pub const NUM_RANDOMS: i64 = 20001; // Conversion / printing functions pub fn convert_to_bigint(x: Value) -> BigUint { let mut big = Default::default(); x.map(|x| { big = BigUint::from_bytes_le(x.to_repr().as_ref()); }); big } pub fn convert_pos_int(x: Value) -> i128 { let bias = 1 << 60; let x_pos = x + Value::known(F::from(bias as u64)); let mut outp: i128 = 0; x_pos.map(|x| { let x_pos = convert_to_u128(&x); let tmp = x_pos as i128 - bias; outp = tmp; }); return outp; } pub fn print_pos_int(prefix: &str, x: Value, scale_factor: u64) { let tmp = convert_pos_int(x); let tmp_float = tmp as f64 / scale_factor as f64; println!("{} x: {} ({})", prefix, tmp, tmp_float); } pub fn print_assigned_arr( prefix: &str, arr: &Vec<&AssignedCell>, scale_factor: u64, ) { for (idx, x) in arr.iter().enumerate() { print_pos_int( &format!("{}[{}]", prefix, idx), x.value().map(|x: &F| x.to_owned()), scale_factor, ); } } // Get the public values pub fn get_public_values() -> Vec { let mut public_vals = vec![]; for val in PUBLIC_VALS.lock().unwrap().iter() { let val = F::from_str_vartime(&val.to_str_radix(10)); public_vals.push(val.unwrap()); } public_vals } // Broadcast fn shape_dominates(s1: &[usize], s2: &[usize]) -> bool { if s1.len() != s2.len() { return false; } for (x1, x2) in s1.iter().zip(s2.iter()) { if x1 < x2 { return false; } } true } // Precondition: s1.len() < s2.len() fn intermediate_shape(s1: &[usize], s2: &[usize]) -> Vec { let mut res = vec![1; s2.len() - s1.len()]; for s in s1.iter() { res.push(*s); } res } fn final_shape(s1: &[usize], s2: &[usize]) -> Vec { let mut res = vec![]; for (x1, x2) in s1.iter().zip(s2.iter()) { res.push(std::cmp::max(*x1, *x2)); } res } pub fn broadcast( x1: &Array, x2: &Array, ) -> (Array, Array) { if x1.shape() == x2.shape() { return (x1.clone(), x2.clone()); } if x1.ndim() == x2.ndim() { let s1 = x1.shape(); let s2 = x2.shape(); if shape_dominates(s1, s2) { return (x1.clone(), x2.broadcast(s1).unwrap().into_owned()); } else if shape_dominates(x2.shape(), x1.shape()) { return (x1.broadcast(s2).unwrap().into_owned(), x2.clone()); } } let (tmp1, tmp2) = if x1.ndim() < x2.ndim() { (x1, x2) } else { (x2, x1) }; // tmp1.ndim() < tmp2.ndim() let s1 = tmp1.shape(); let s2 = tmp2.shape(); let s = intermediate_shape(s1, s2); let final_shape = final_shape(s2, s.as_slice()); let tmp1 = tmp1.broadcast(s.clone()).unwrap().into_owned(); let tmp1 = tmp1.broadcast(final_shape.as_slice()).unwrap().into_owned(); let tmp2 = tmp2.broadcast(final_shape.as_slice()).unwrap().into_owned(); // println!("x1: {:?} x2: {:?}", x1.shape(), x2.shape()); // println!("s1: {:?} s2: {:?} s: {:?}", s1, s2, s); // println!("tmp1 shape: {:?}", tmp1.shape()); // println!("tmp2 shape: {:?}", tmp2.shape()); if x1.ndim() < x2.ndim() { return (tmp1, tmp2); } else { return (tmp2, tmp1); } } ================================================ FILE: src/utils/loader.rs ================================================ use std::{fs::File, io::BufReader}; use serde_derive::{Deserialize, Serialize}; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct TensorMsgpack { pub idx: i64, pub shape: Vec, pub data: Vec, } #[derive(Clone, Debug, Serialize, Deserialize)] pub struct LayerMsgpack { pub layer_type: String, pub params: Vec, pub inp_idxes: Vec, pub inp_shapes: Vec>, pub out_idxes: Vec, pub out_shapes: Vec>, pub mask: Vec, } #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ModelMsgpack { pub global_sf: i64, pub k: i64, pub num_cols: i64, pub inp_idxes: Vec, pub out_idxes: Vec, pub tensors: Vec, pub layers: Vec, pub use_selectors: Option, pub commit_before: Option>>, pub commit_after: Option>>, pub bits_per_elem: Option, // Specifically for packing for the commitments pub num_random: Option, } pub fn load_config_msgpack(config_path: &str) -> ModelMsgpack { let model: ModelMsgpack = { let file = File::open(config_path).unwrap(); let mut reader = BufReader::new(file); rmp_serde::from_read(&mut reader).unwrap() }; model } pub fn load_model_msgpack(config_path: &str, inp_path: &str) -> ModelMsgpack { let mut model = load_config_msgpack(config_path); let inp: Vec = { let file = File::open(inp_path).unwrap(); let mut reader = BufReader::new(file); rmp_serde::from_read(&mut reader).unwrap() }; for tensor in inp { model.tensors.push(tensor); } // Default to using selectors, commit if use_selectors is not specified if model.use_selectors.is_none() { model.use_selectors = Some(true) }; if model.commit_before.is_none() { model.commit_before = Some(vec![]) }; if model.commit_after.is_none() { model.commit_after = Some(vec![]) }; if model.bits_per_elem.is_none() { model.bits_per_elem = Some(model.k) }; if model.num_random.is_none() { model.num_random = Some(20001) }; model } ================================================ FILE: src/utils/proving_ipa.rs ================================================ use std::{ fs::File, io::{BufReader, Write}, path::Path, time::Instant, }; use halo2_proofs::{ dev::MockProver, halo2curves::pasta::{EqAffine, Fp}, plonk::{create_proof, keygen_pk, keygen_vk, verify_proof}, poly::{ commitment::{Params, ParamsProver}, ipa::{ commitment::{IPACommitmentScheme, ParamsIPA}, multiopen::ProverIPA, strategy::SingleStrategy, }, VerificationStrategy, }, transcript::{ Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer, }, }; use crate::{model::ModelCircuit, utils::helpers::get_public_values}; pub fn get_ipa_params(params_dir: &str, degree: u32) -> ParamsIPA { let path = format!("{}/{}.params", params_dir, degree); let params_path = Path::new(&path); if File::open(¶ms_path).is_err() { let params: ParamsIPA = ParamsIPA::new(degree); let mut buf = Vec::new(); params.write(&mut buf).expect("Failed to write params"); let mut file = File::create(¶ms_path).expect("Failed to create params file"); file .write_all(&buf[..]) .expect("Failed to write params to file"); } let params_fs = File::open(¶ms_path).expect("couldn't load params"); let params: ParamsIPA = Params::read::<_>(&mut BufReader::new(params_fs)).expect("Failed to read params"); params } pub fn time_circuit_ipa(circuit: ModelCircuit) { let rng = rand::thread_rng(); let start = Instant::now(); let degree = circuit.k as u32; let empty_circuit = circuit.clone(); let proof_circuit = circuit; let params = get_ipa_params("./params_ipa", degree); let circuit_duration = start.elapsed(); println!( "Time elapsed in params construction: {:?}", circuit_duration ); let vk = keygen_vk(¶ms, &empty_circuit).unwrap(); let vk_duration = start.elapsed(); println!( "Time elapsed in generating vkey: {:?}", vk_duration - circuit_duration ); let pk = keygen_pk(¶ms, vk, &empty_circuit).unwrap(); let pk_duration = start.elapsed(); println!( "Time elapsed in generating pkey: {:?}", pk_duration - vk_duration ); drop(empty_circuit); let fill_duration = start.elapsed(); let _prover = MockProver::run(degree, &proof_circuit, vec![vec![]]).unwrap(); let public_vals = get_public_values(); println!( "Time elapsed in filling circuit: {:?}", fill_duration - pk_duration ); let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]); create_proof::, ProverIPA, _, _, _, _>( ¶ms, &pk, &[proof_circuit], &[&[&public_vals]], rng, &mut transcript, ) .unwrap(); let proof = transcript.finalize(); let proof_duration = start.elapsed(); println!("Proving time: {:?}", proof_duration - fill_duration); let proof_size = { let mut folder = std::path::PathBuf::new(); folder.push("proof"); let mut fd = std::fs::File::create(folder.as_path()).unwrap(); folder.pop(); fd.write_all(&proof).unwrap(); fd.metadata().unwrap().len() }; println!("Proof size: {} bytes", proof_size); let strategy = SingleStrategy::new(¶ms); let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]); assert!( verify_proof( ¶ms, pk.get_vk(), strategy, &[&[&public_vals]], &mut transcript ) .is_ok(), "proof did not verify" ); let verify_duration = start.elapsed(); println!("Verifying time: {:?}", verify_duration - proof_duration); } ================================================ FILE: src/utils/proving_kzg.rs ================================================ use std::{ fs::File, io::{BufReader, Write}, path::Path, time::Instant, }; use halo2_proofs::{ dev::MockProver, halo2curves::bn256::{Bn256, Fr, G1Affine}, plonk::{create_proof, keygen_pk, keygen_vk, verify_proof, VerifyingKey}, poly::{ commitment::Params, kzg::{ commitment::{KZGCommitmentScheme, ParamsKZG}, multiopen::{ProverSHPLONK, VerifierSHPLONK}, strategy::SingleStrategy, }, }, transcript::{ Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer, }, SerdeFormat, }; use crate::{model::ModelCircuit, utils::helpers::get_public_values}; pub fn get_kzg_params(params_dir: &str, degree: u32) -> ParamsKZG { let rng = rand::thread_rng(); let path = format!("{}/{}.params", params_dir, degree); let params_path = Path::new(&path); if File::open(¶ms_path).is_err() { let params = ParamsKZG::::setup(degree, rng); let mut buf = Vec::new(); params.write(&mut buf).expect("Failed to write params"); let mut file = File::create(¶ms_path).expect("Failed to create params file"); file .write_all(&buf[..]) .expect("Failed to write params to file"); } let mut params_fs = File::open(¶ms_path).expect("couldn't load params"); let params = ParamsKZG::::read(&mut params_fs).expect("Failed to read params"); params } pub fn serialize(data: &Vec, path: &str) -> u64 { let mut file = File::create(path).unwrap(); file.write_all(data).unwrap(); file.metadata().unwrap().len() } pub fn verify_kzg( params: &ParamsKZG, vk: &VerifyingKey, strategy: SingleStrategy, public_vals: &Vec, mut transcript: Blake2bRead<&[u8], G1Affine, Challenge255>, ) { assert!( verify_proof::< KZGCommitmentScheme, VerifierSHPLONK<'_, Bn256>, Challenge255, Blake2bRead<&[u8], G1Affine, Challenge255>, halo2_proofs::poly::kzg::strategy::SingleStrategy<'_, Bn256>, >(¶ms, &vk, strategy, &[&[&public_vals]], &mut transcript) .is_ok(), "proof did not verify" ); } pub fn time_circuit_kzg(circuit: ModelCircuit) { let rng = rand::thread_rng(); let start = Instant::now(); let degree = circuit.k as u32; let params = get_kzg_params("./params_kzg", degree); let circuit_duration = start.elapsed(); println!( "Time elapsed in params construction: {:?}", circuit_duration ); let vk_circuit = circuit.clone(); let vk = keygen_vk(¶ms, &vk_circuit).unwrap(); drop(vk_circuit); let vk_duration = start.elapsed(); println!( "Time elapsed in generating vkey: {:?}", vk_duration - circuit_duration ); let vkey_size = serialize(&vk.to_bytes(SerdeFormat::RawBytes), "vkey"); println!("vkey size: {} bytes", vkey_size); let pk_circuit = circuit.clone(); let pk = keygen_pk(¶ms, vk, &pk_circuit).unwrap(); let pk_duration = start.elapsed(); println!( "Time elapsed in generating pkey: {:?}", pk_duration - vk_duration ); drop(pk_circuit); let pkey_size = serialize(&pk.to_bytes(SerdeFormat::RawBytes), "pkey"); println!("pkey size: {} bytes", pkey_size); let fill_duration = start.elapsed(); let proof_circuit = circuit.clone(); let _prover = MockProver::run(degree, &proof_circuit, vec![vec![]]).unwrap(); let public_vals = get_public_values(); println!( "Time elapsed in filling circuit: {:?}", fill_duration - pk_duration ); // Convert public vals to serializable format let public_vals_u8: Vec = public_vals .iter() .map(|v: &Fr| v.to_bytes().to_vec()) .flatten() .collect(); let public_vals_u8_size = serialize(&public_vals_u8, "public_vals"); println!("Public vals size: {} bytes", public_vals_u8_size); let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]); create_proof::< KZGCommitmentScheme, ProverSHPLONK<'_, Bn256>, Challenge255, _, Blake2bWrite, G1Affine, Challenge255>, ModelCircuit, >( ¶ms, &pk, &[proof_circuit], &[&[&public_vals]], rng, &mut transcript, ) .unwrap(); let proof = transcript.finalize(); let proof_duration = start.elapsed(); println!("Proving time: {:?}", proof_duration - fill_duration); let proof_size = serialize(&proof, "proof"); let proof = std::fs::read("proof").unwrap(); println!("Proof size: {} bytes", proof_size); let strategy = SingleStrategy::new(¶ms); let transcript_read = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]); println!("public vals: {:?}", public_vals); verify_kzg( ¶ms, &pk.get_vk(), strategy, &public_vals, transcript_read, ); let verify_duration = start.elapsed(); println!("Verifying time: {:?}", verify_duration - proof_duration); } // Standalone verification pub fn verify_circuit_kzg( circuit: ModelCircuit, vkey_fname: &str, proof_fname: &str, public_vals_fname: &str, ) { let degree = circuit.k as u32; let params = get_kzg_params("./params_kzg", degree); println!("Loaded the parameters"); let vk = VerifyingKey::read::, ModelCircuit>( &mut BufReader::new(File::open(vkey_fname).unwrap()), SerdeFormat::RawBytes, (), ) .unwrap(); println!("Loaded vkey"); let proof = std::fs::read(proof_fname).unwrap(); let public_vals_u8 = std::fs::read(&public_vals_fname).unwrap(); let public_vals: Vec = public_vals_u8 .chunks(32) .map(|chunk| Fr::from_bytes(chunk.try_into().expect("conversion failed")).unwrap()) .collect(); let strategy = SingleStrategy::new(¶ms); let transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]); let start = Instant::now(); let verify_start = start.elapsed(); verify_kzg(¶ms, &vk, strategy, &public_vals, transcript); let verify_duration = start.elapsed(); println!("Verifying time: {:?}", verify_duration - verify_start); println!("Proof verified!") } ================================================ FILE: src/utils.rs ================================================ pub mod helpers; pub mod loader; pub mod proving_ipa; pub mod proving_kzg; ================================================ FILE: testing/circuits/last_two_layers.py ================================================ import tensorflow as tf import os import numpy as np interpreter = tf.lite.Interpreter( model_path=f'./testing/circuits/v2_1.0_224.tflite' ) interpreter.allocate_tensors() NAME_TO_TENSOR = {} for tensor_details in interpreter.get_tensor_details(): NAME_TO_TENSOR[tensor_details['name']] = tensor_details Wc = interpreter.get_tensor(NAME_TO_TENSOR['Const_71']['index']) Bc = interpreter.get_tensor(NAME_TO_TENSOR['MobilenetV2/Conv_1/Conv2D_bias']['index']) class LastTwoLayers(tf.keras.Model): def __init__(self, name=None): super().__init__(name = name) self.a_variable = tf.Variable(5.0, name="train_me") self.conv1 = tf.keras.layers.Conv2D( 1280, (1, 1), activation='relu6', padding='same', input_shape=(1, 7, 7, 320) ) self.avg_pool = tf.keras.layers.AveragePooling2D( pool_size=(7, 7), padding='valid', strides=(1, 1) ) self.conv2 = tf.keras.layers.Conv2D( 102, (1, 1), padding='valid', input_shape=(1, 1, 1, 1280) ) self.softmax = tf.keras.layers.Softmax() def call(self, x): x = self.conv1(x) x = self.avg_pool(x) x = self.conv2(x) x = tf.reshape(x, [1, 102]) x = self.softmax(x) return x my_sequential_model = LastTwoLayers(name="the_model") my_sequential_model.compile(optimizer='sgd', loss='categorical_crossentropy') x = np.random.random((1, 7, 7, 320)) my_sequential_model.predict(x) my_sequential_model.conv1.set_weights([np.transpose(Wc, [1,2,3,0]), Bc]) # x, y, chan_in, chan_out # 1 batch, 7 height, 320 width, 7 channels # 7 height, 1 width, 7 channels, 1280 out channels # 1 height, 320 width, 1280 out channels # 1 Batch, 7 height, 7 width, 320 channels () [This is the input to the layer] # 1 Batch, 7 height, 7 width, 1280 channels (dout) [This is the output of another layer] # We want to transform this so that we rotate the input by 90 degrees W = np.zeros([1, 1, 1280, 102]) my_sequential_model.conv2.set_weights([ W, np.zeros([102]) ]) converter = tf.lite.TFLiteConverter.from_keras_model(my_sequential_model) tflite_model = converter.convert() with open('./examples/v2_1.0_224_truncated/v2_1.0_224_truncated.tflite', 'wb') as f: f.write(tflite_model) ================================================ FILE: testing/circuits/v2_1.0_224.tflite ================================================ [File too large to display: 13.3 MB]