This project expands the functionality of the zkLLVM compiler by integrating an ONNX frontend, leveraging the capabilities of the ONNX-MLIR compiler technology. It allows developers to compile and assign ZK circuits based on their pre-trained ONNX models and validate them using the Placeholder proof system (refer to the zkLLVM documentation for more details).
This project includes two binaries:
zkml-onnx-compiler
: A compiler designed to lower .onnx files. (Note: The name of this compiler may be subject to change as the project evolves.)mlir-assigner
An assigner responsible for interpreting .mlir dialect files and generating a plonkish circuit along with an assigned table.
We provide a comprehensive list of the ONNX operations currently supported by our system, along with any limitations that may apply.
We heavily encourage to use the nix
devshell when working with, or building
from source. It is possible to build the project on Linux (tested with Ubuntu
22.04), but the build process can be involved.
Included in the project is a Nix flake that facilitates setting up a development
shell with all necessary dependencies for building the mlir-assigner
. As part
of this setup process, the flake fetches and builds the bundled LLVM within
zkLLVM. As this is can be pretty time-consuming, we recommend passing the
--cores 16
flag or a similar option during the initial invocation of
nix develop
.
nix develop # or
nix develop --cores 16 #allow the first-time build to use more cores for building the deps
Once inside the Nix shell environment, proceed to the following steps:
mkdir build && cd build
cmake -DMLIR_DIR=${MLIR_DIR} ..
make -j
The resulting binaries can be located in the build/bin
directory.
For an optimal development environment setup, we strongly recommend utilizing
the Nix
devshell. However, if you prefer to build on your native machine,
please follow the steps outlined below.
- Download the source
git clone --recurse-submodules -j8 [email protected]:TaceoLabs/zkllvm-mlir-assigner.git && cd zkllvm-mlir-assigner
-
One of our dependencies is ONNX-MLIR (
libs/onnx-mlir
). Ensure that all requirements and build instructions for ONNX-MLIR are correctly set up. Although you don't need to build it explicitly, everything must be configured correctly. This includes having a build tree of MLIR (not ONNX-MLIR) as part of the LLVM project, as outlined here. -
Confirm that your environment variable
$MLIR_DIR
points to an MLIR CMake module by running:echo $MLIR_DIR
-
Setup cmake and build
mkdir build && cd build cmake -DMLIR_DIR=${MLIR_DIR} .. make -j zkml-onnx-compiler mlir-assigner
This section shows how to use the zkML frontend for zkLLVM. In this example, we guide you through every step to proof the CNN-MNIST Model from the ONNX model zoo on Ubuntu 22.04.
We expect that you already built the two binaries build/bin/zk-ml-opt
and
build/bin/mlir-assigner
from source or obtained them in another way. If not,
follow the build instructions.
-
Setup: We start by creating an empty folder where we place our binaries and our model. So when you built from source we do:
mkdir CNN-Mnist cp build/bin/mlir-assigner CNN-Mnist && cp build/bin/zkml-onnx-compiler CNN-Mnist cd CNN-Mnist wget https://github.com/onnx/models/raw/main/validated/vision/classification/mnist/model/mnist-12.onnx
In case you have another model you want to proof, use that instead of the CNN-Mnist model.
-
Compile ONNX to MLIR: Having your pre-trained model at place, we use the
zkml-onnx-compiler
to compile the model to.mlir
.When lowering to
.mlir
you also have to specify which inputs of the model are public/private by passing a comma seperated list to thezk
flag (e.g.--zk 1,1,0
). If an input is set to private, the produced Zero-Knowledge proof will not reveal the input.For instance, imagine we want to keep the single input to the mnist model we use (the picture of the digit) private. To achieve this, we call the compiler with the following command
./zkml-onnx-compiler mnist-12.onnx --output mnist-12.mlir --zk 1
The
zkml-onnx-compiler
can also lower the model to different IRs. Have a look by adding the--help
flag.The emitted
.mlir
consists of the dialects defined by ONNX-MLIR and an additional dialect defined by this project with the namespacezkML
. This dialect defines operations which need special handling in contrast to the default lowering provided by ONNX-MLIR, e.g., as seen during the lowering of MatMuls, Gemms, or Convolutions. Natively lowering matrix multiplications leads to polluted traces with a lot of additions and multiplications. For that we introduced the operationzkml.dot-product
which improves performance drastically. -
Prepare your input: Before running the
mlir-assigner
, ensure that your input is properly formatted in JSON format. The input should be separated into two files: one for the public inputs and one for the private inputs. Each input file should consist of an array of objects representing tensors, denoted asmemref
.Each memref object should contain the following four elements:
- idx: The index in the input tensor (public and private have separate indices).
- data: The flattened data of the input tensor.
- dims: The dimensions of the input tensor.
- type: The type of the elements of the tensor.
Have a look at e.g.,the the input to this small model to see how you should prepare your input. Additionally, we prepared an json schema you can use as reference.
You can refer to an example input file, such as the one provided for the CNN-Mnist model, to see how the input should be prepared. Additionally, we have prepared a JSON schema that you can use as a reference.
For example, to use the prepared input file for the CNN-Mnist model, you can simply copy it to your current folder:
cp ../mlir-assigner/tests/Models/ConvMnist/mnist-12.json .
-
Perform inference and assign circuit: The next step involves performing inference of the model using the provided input within the plonkish arithmetization of the proof system. This process may take some time, depending on the complexity of your model.
To perform this step, execute the following command:
./mlir-assigner -b mnist-12.mlir -p mnist-12.json -e pallas -c circuit.crt -o mnist-12.output.json -t assignment.tbl -f dec --check
Again, have a look on the configurations of the
mlir-assigner
by adding the--help
flag.Remember that we want to keep the input to the model private, therefore we have to submit the private input with the
-p
flag. The public input , if applicable, corresponds to the-i
flag, omitted in this example. You can find the unassigned circuit in the filecircuit.crt
and the assignment inassignment.tbl
, as you are used to fromzkLLVM
. We refer to the documentation from zkLLVM on how to produce proofs withPlaceholder
and/or publish your circuits to the proof market.
Certain limitations arise from the fact that we are computing Zero-Knowledge proofs over the execution of models. First, and foremost, as we are building on the ONNX-MLIR compiler technology. Therefore, we are limited to the ONNX operations supported by ONNX-MLIR. Additionally, certain operations are not possible to depict in Zero-Knowledge. This mostly includes operations which would lead to a "dynamic" circuit, meaning, a circuit that changes depending on the input. For a comprehensive list of the supported ONNX operations, have a look in the test folder.
Additionally, ONNX supports dynamic-sized input dimensions, commonly used for
parameters like batch_size
to group evaluations for a single run of the model.
However, we do not support dynamic inputs, as we require all tensor shapes to be
publicly known. This assumption is essential for constructing meaningful and
sound proofs. Consequently, the batch dimension must be fixed at compile time.
To check if your model has dynamic inputs, we recommend the Netron App. It's an excellent tool for visualizing and inspecting ONNX models. Below is an example illustrating a dynamic input:
The input highlights the dynamic input dimension
batch_size
. This is a screenshot of the nano-GPT model.
To solve this problem, follow the steps outlined in the ONNX docs.
To test your build, have a look in the test folder.
This is experimental software and is provided on an "as is" and "as available" basis. We do not give any warranties and will not be liable for any losses incurred through any use of this code base.