Skip to content

Commit

Permalink
refactor: Op traits (#194)
Browse files Browse the repository at this point in the history
  • Loading branch information
@alexander-camuto
alexander-camuto authored Apr 20, 2023
1 parent a481b04 commit 0796594
Show file tree
Hide file tree
Showing 40 changed files with 2,908 additions and 2,765 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ tokio = { version = "1.26.0", features = ["macros", "rt"] }
puruspe = "0.2.0"
bincode = "*"


# python binding related deps
pyo3 = { version = "0.18.2", features = ["extension-module", "abi3-py37"], optional = true }
pyo3-log = { version = "0.8.1", optional = true }
Expand Down
8 changes: 7 additions & 1 deletion benches/accum_affine.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -53,7 +54,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Affine.into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Affine),
)
.unwrap();
Ok(())
},
Expand Down
8 changes: 4 additions & 4 deletions benches/accum_conv.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -61,14 +62,13 @@ impl Circuit<Fr> for MyCircuit {
|mut region| {
config
.layout(
&mut region,
Some(&mut region),
&[self.image.clone(), self.kernel.clone(), self.bias.clone()],
&mut 0,
Op::Conv {
Box::new(PolyOp::Conv {
padding: (0, 0),
stride: (1, 1),
}
.into(),
}),
)
.unwrap();
Ok(())
Expand Down
8 changes: 7 additions & 1 deletion benches/accum_dot.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -51,7 +52,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Dot.into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Dot),
)
.unwrap();
Ok(())
},
Expand Down
8 changes: 7 additions & 1 deletion benches/accum_matmul.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -53,7 +54,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Matmul.into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Matmul),
)
.unwrap();
Ok(())
},
Expand Down
10 changes: 6 additions & 4 deletions benches/accum_matmul_relu.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::*;

use ezkl_lib::circuit::lookup::LookupOp;
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
use ezkl_lib::pfsys::{create_keys, gen_srs};
Expand Down Expand Up @@ -65,19 +67,19 @@ impl Circuit<Fr> for MyCircuit {
layouter.assign_region(
|| "",
|mut region| {
let op = Op::Matmul;
let op = PolyOp::Matmul;
let mut offset = 0;
let output = config
.base_config
.layout(&mut region, &self.inputs, &mut offset, op.into())
.layout(Some(&mut region), &self.inputs, &mut offset, Box::new(op))
.unwrap();
let _output = config
.base_config
.layout(
&mut region,
Some(&mut region),
&[output.unwrap()],
&mut offset,
LookupOp::ReLU { scale: 1 }.into(),
Box::new(LookupOp::ReLU { scale: 1 }),
)
.unwrap();
Ok(())
Expand Down
8 changes: 7 additions & 1 deletion benches/accum_pack.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -51,7 +52,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Pack(2, 1).into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Pack(2, 1)),
)
.unwrap();
Ok(())
},
Expand Down
8 changes: 7 additions & 1 deletion benches/accum_sum.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -51,7 +52,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Sum.into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Sum),
)
.unwrap();
Ok(())
},
Expand Down
8 changes: 4 additions & 4 deletions benches/accum_sumpool.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -56,15 +57,14 @@ impl Circuit<Fr> for MyCircuit {
|mut region| {
config
.layout(
&mut region,
Some(&mut region),
&[self.image.clone()],
&mut 0,
Op::SumPool {
Box::new(PolyOp::SumPool {
padding: (0, 0),
stride: (1, 1),
kernel_shape: (2, 2),
}
.into(),
}),
)
.unwrap();
Ok(())
Expand Down
8 changes: 7 additions & 1 deletion benches/pairwise_add.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -51,7 +52,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Add.into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Add),
)
.unwrap();
Ok(())
},
Expand Down
8 changes: 7 additions & 1 deletion benches/pairwise_pow.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::poly::PolyOp;
use ezkl_lib::circuit::*;
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
Expand Down Expand Up @@ -51,7 +52,12 @@ impl Circuit<Fr> for MyCircuit {
|| "",
|mut region| {
config
.layout(&mut region, &self.inputs, &mut 0, Op::Pow(4).into())
.layout(
Some(&mut region),
&self.inputs,
&mut 0,
Box::new(PolyOp::Pow(4)),
)
.unwrap();
Ok(())
},
Expand Down
6 changes: 3 additions & 3 deletions benches/relu.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use ezkl_lib::circuit::{BaseConfig as Config, CheckMode, LookupOp};
use ezkl_lib::circuit::{ops::lookup::LookupOp, BaseConfig as Config, CheckMode};
use ezkl_lib::commands::TranscriptType;
use ezkl_lib::execute::create_proof_circuit_kzg;
use ezkl_lib::pfsys::{create_keys, gen_srs};
Expand Down Expand Up @@ -59,10 +59,10 @@ impl Circuit<Fr> for NLCircuit {
|mut region| {
config
.layout(
&mut region,
Some(&mut region),
&[self.input.clone()],
&mut 0,
LookupOp::ReLU { scale: 128 }.into(),
Box::new(LookupOp::ReLU { scale: 128 }),
)
.unwrap();
Ok(())
Expand Down
16 changes: 9 additions & 7 deletions examples/conv2d_mnist/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use ezkl_lib::circuit::{BaseConfig as PolyConfig, CheckMode, LookupOp, Op as PolyOp};
use ezkl_lib::circuit::{
ops::lookup::LookupOp, ops::poly::PolyOp, BaseConfig as PolyConfig, CheckMode,
};
use ezkl_lib::fieldutils;
use ezkl_lib::fieldutils::i32_to_felt;
use ezkl_lib::tensor::*;
Expand Down Expand Up @@ -178,35 +180,35 @@ where
let x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[
self.input.clone(),
self.l0_params[0].clone(),
self.l0_params[1].clone(),
],
&mut offset,
op.into(),
Box::new(op),
)
.unwrap();

let mut x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x.unwrap()],
&mut offset,
LookupOp::ReLU { scale: 32 }.into(),
Box::new(LookupOp::ReLU { scale: 32 }),
)
.unwrap()
.unwrap();
x.flatten();
let l2out = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x, self.l2_params[0].clone(), self.l2_params[1].clone()],
&mut offset,
PolyOp::Affine.into(),
Box::new(PolyOp::Affine),
)
.unwrap();
Ok(l2out)
Expand Down
27 changes: 14 additions & 13 deletions examples/mlp_4d.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
use ezkl_lib::circuit::{BaseConfig as PolyConfig, CheckMode, LookupOp, Op as PolyOp};
use ezkl_lib::circuit::{
ops::lookup::LookupOp, ops::poly::PolyOp, BaseConfig as PolyConfig, CheckMode,
};
use ezkl_lib::fieldutils::i32_to_felt;
use ezkl_lib::tensor::*;
use halo2_proofs::dev::MockProver;
Expand Down Expand Up @@ -95,14 +97,14 @@ impl<F: FieldExt + TensorType, const LEN: usize, const BITS: usize> Circuit<F>
let x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[
self.input.clone(),
self.l0_params[0].clone(),
self.l0_params[1].clone(),
],
&mut offset,
PolyOp::Affine.into(),
Box::new(PolyOp::Affine),
)
.unwrap()
.unwrap();
Expand All @@ -112,10 +114,10 @@ impl<F: FieldExt + TensorType, const LEN: usize, const BITS: usize> Circuit<F>
let mut x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x],
&mut offset,
LookupOp::ReLU { scale: 1 }.into(),
Box::new(LookupOp::ReLU { scale: 1 }),
)
.unwrap()
.unwrap();
Expand All @@ -125,10 +127,10 @@ impl<F: FieldExt + TensorType, const LEN: usize, const BITS: usize> Circuit<F>
let x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x, self.l2_params[0].clone(), self.l2_params[1].clone()],
&mut offset,
PolyOp::Affine.into(),
Box::new(PolyOp::Affine),
)
.unwrap()
.unwrap();
Expand All @@ -137,23 +139,22 @@ impl<F: FieldExt + TensorType, const LEN: usize, const BITS: usize> Circuit<F>
let x = config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x],
&mut offset,
LookupOp::ReLU { scale: 1 }.into(),
Box::new(LookupOp::ReLU { scale: 1 }),
)
.unwrap();
println!("offset: {}", offset);
Ok(config
.layer_config
.layout(
&mut region,
Some(&mut region),
&[x.unwrap()],
&mut offset,
LookupOp::Div {
Box::new(LookupOp::Div {
denom: ezkl_lib::circuit::utils::F32::from(128.),
}
.into(),
}),
)
.unwrap())
},
Expand Down
18 changes: 18 additions & 0 deletions examples/onnx/1l_instance_norm/gen.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from torch import nn
from ezkl import export
import torch

class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()

self.layer = nn.InstanceNorm2d(3).eval()

def forward(self, x):
return [self.layer(x)]

circuit = MyModel()
export(circuit, input_shape = [3,2,2])



Loading

0 comments on commit 0796594

Please sign in to comment.