Skip to content

TaceoLabs/zkllvm-mlir-assigner

This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Folders and files

NameName
Last commit message
Last commit date

Latest commit

3f2274a · Jan 9, 2024

History

76 Commits
Dec 6, 2023
Dec 5, 2023
Dec 6, 2023
Jan 9, 2024
Jan 9, 2024
Dec 20, 2023
Dec 20, 2023
Dec 18, 2023
Dec 1, 2023
Dec 5, 2023
Dec 1, 2023
Dec 5, 2023
Nov 3, 2023
Jan 5, 2024
Nov 21, 2023
Dec 18, 2023
Jan 5, 2024
Dec 18, 2023

Repository files navigation

ONNX zkML Frontend for zkLLVM

X (formerly Twitter) Follow

This project extends the zkLLVM compiler with an ONNX frontend, leveraging the ONNX-MLIR compiler technology. It allows developers to compile and assign ZK circuits from their pre-trained onnx models and proof them with the Placeholder proof system (see zkLLVM documentation).

This project builds two binaries, namely:

  • zkml-onnx-compiler A compiler that lowers .onnx files. (name is subject to change)
  • mlir-assigner An assigner that interprets .mlir dialect files and produces a plonkish circuit and an assigned table.

We list the currently supported ONNX operations and limitations here.

Build

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.

Build in nix devshell

The included nix flake sets up a devshell with all dependencies for building the mlir-assigner. Since it also fetches and builds the bundled LLVM in zkLLVM, this may take some time, so we recommend to pass --cores 16 or similar to 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
# in the devshell
mkdir build && cd build
cmake -DMLIR_DIR=${MLIR_DIR} ..
make -j

You can find the binaries in the build/bin folder.

Build on Ubuntu 22.04

We stress again to use the preferred approach, the nix devshell. If you still want to built on your native machine, follow these steps.

  1. Download the source:
git clone --recurse-submodules -j8 [email protected]:TaceoLabs/zkllvm-mlir-assigner.git && cd zkllvm-mlir-assigner
  1. One of our dependencies is ONNX-MLIR (libs/onnx-mlir). Here you can find the requirements and built instructions for ONNX-MLIR. You do not have to built it, but every thing needs to be setup correctly. This includes a build tree of MLIR (not ONNX-MLIR) as part of the LLVM project, as seen here.

  2. At this point your environment variable MLIR_DIR should point to a MLIR cmake module. Verify that by writing

echo $MLIR_DIR
  1. Setup cmake and build
mkdir build && cd build
cmake -DMLIR_DIR=${MLIR_DIR} ..
make -j zkml-onnx-compiler mlir-assigner

Testing

To test your build, have a look in the test folder.

Example

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.

  1. Setup: We start by creating an empty folder here 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/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.

  1. Compile ONNX to MLIR: Having your pre-trained model at place, we use the zkml-onnx-compiler to compile the model to .mlir. compile You can do this by calling the zkml-onnx-compiler:
./zkml-onnx-compiler mnist-12.onnx -i mnist-12.mlir

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 namespace zkML. 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 operation zkml.Dot-Product which improves performance drastically.

  1. Prepare your input: The mlir-assigner expects the input for inference in json format. Have a look at e.g., the input to this small model to see how you should prepare your input. For this example, we have an input file prepared. Therefore, we just copy it to our folder:
cp ../mlir-assigner/tests/Models/ConvMnist/mnist-12.json .

inference 4. Perform inference and assign circuit: The next step may take some time depending on your model. We have to perform inference of the model and the provided input within the plonkish arithmetization of the proof system. We do this by calling:

./mlir-assigner -b mnist-12.mlir -i mnist-12.json -e pallas -c circuit.crt -t assignment.tbl --print_circuit_output --check

Again, have a look on the configurations of the mlir-assigner by adding the --help flag.

You can find the unassigned circuit in the file circuit.crt and the assignment in assignment.tbl, as you are used to from zkLLVM. We refer to the documentation from zkLLVM on how to produce proofs with Placeholder and/or publish your circuits to the proof market.

Disclaimer

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.