Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(zkvm): add cycle tracker print to precompilees #8

Merged
merged 4 commits into from
Jun 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file added .cargo-ok
Empty file.
2 changes: 2 additions & 0 deletions crates/precompile/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ portable = ["revm-primitives/portable", "c-kzg?/portable"]
# In Linux it passes. If you don't require to build wasm on win/mac, it is safe to use it and it is enabled by default.
secp256k1 = ["dep:secp256k1"]

sp1-cycle-tracker = []

[[bench]]
name = "bench"
path = "benches/bench.rs"
Expand Down
49 changes: 28 additions & 21 deletions crates/precompile/src/blake2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ pub const FUN: PrecompileWithAddress =
/// input format:
/// [4 bytes for rounds][64 bytes for h][128 bytes for m][8 bytes for t_0][8 bytes for t_1][1 byte for f]
pub fn run(input: &Bytes, gas_limit: u64) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: blake2");

let input = &input[..];

if input.len() != INPUT_LENGTH {
Expand All @@ -32,33 +35,37 @@ pub fn run(input: &Bytes, gas_limit: u64) -> PrecompileResult {
}

let out = if zk_op::contains_operation(&ZkOperation::Blake2) {
zk_op::ZKVM_OPERATOR.get().unwrap().blake2_run(input).unwrap()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.blake2_run(input)
.unwrap()
} else {
// Not indented to keep the diff clean and make changes to the original code obvious

let mut h = [0u64; 8];
let mut m = [0u64; 16];
let mut h = [0u64; 8];
let mut m = [0u64; 16];

for (i, pos) in (4..68).step_by(8).enumerate() {
h[i] = u64::from_le_bytes(input[pos..pos + 8].try_into().unwrap());
}
for (i, pos) in (68..196).step_by(8).enumerate() {
m[i] = u64::from_le_bytes(input[pos..pos + 8].try_into().unwrap());
}
let t = [
u64::from_le_bytes(input[196..196 + 8].try_into().unwrap()),
u64::from_le_bytes(input[204..204 + 8].try_into().unwrap()),
];
for (i, pos) in (4..68).step_by(8).enumerate() {
h[i] = u64::from_le_bytes(input[pos..pos + 8].try_into().unwrap());
}
for (i, pos) in (68..196).step_by(8).enumerate() {
m[i] = u64::from_le_bytes(input[pos..pos + 8].try_into().unwrap());
}
let t = [
u64::from_le_bytes(input[196..196 + 8].try_into().unwrap()),
u64::from_le_bytes(input[204..204 + 8].try_into().unwrap()),
];

algo::compress(rounds, &mut h, m, t, f);
algo::compress(rounds, &mut h, m, t, f);

let mut out = [0u8; 64];
for (i, h) in (0..64).step_by(8).zip(h.iter()) {
out[i..i + 8].copy_from_slice(&h.to_le_bytes());
}
let mut out = [0u8; 64];
for (i, h) in (0..64).step_by(8).zip(h.iter()) {
out[i..i + 8].copy_from_slice(&h.to_le_bytes());
}

out
out
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: blake2");

Ok((gas_used, out.into()))
}
Expand Down
157 changes: 88 additions & 69 deletions crates/precompile/src/bn128.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,57 +122,68 @@ pub fn new_g1_point(px: Fq, py: Fq) -> Result<G1, Error> {
}

pub fn run_add(input: &[u8], gas_cost: u64, gas_limit: u64) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: bn-add");
if gas_cost > gas_limit {
return Err(Error::OutOfGas);
}

let output = if zk_op::contains_operation(&ZkOperation::Bn128Add) {
zk_op::ZKVM_OPERATOR.get().unwrap().bn128_run_add(input).unwrap()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.bn128_run_add(input)
.unwrap()
} else {
// Not indented to keep the diff clean and make changes to the original code obvious
let input = right_pad::<ADD_INPUT_LEN>(input);

let input = right_pad::<ADD_INPUT_LEN>(input);
let p1 = read_point(&input[..64])?;
let p2 = read_point(&input[64..])?;

let p1 = read_point(&input[..64])?;
let p2 = read_point(&input[64..])?;

let mut output = [0u8; 64];
if let Some(sum) = AffineG1::from_jacobian(p1 + p2) {
sum.x().to_big_endian(&mut output[..32]).unwrap();
sum.y().to_big_endian(&mut output[32..]).unwrap();
}
output
let mut output = [0u8; 64];
if let Some(sum) = AffineG1::from_jacobian(p1 + p2) {
sum.x().to_big_endian(&mut output[..32]).unwrap();
sum.y().to_big_endian(&mut output[32..]).unwrap();
}

output
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: bn-add");

Ok((gas_cost, output.into()))
}

pub fn run_mul(input: &[u8], gas_cost: u64, gas_limit: u64) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: bn-mul");
if gas_cost > gas_limit {
return Err(Error::OutOfGas);
}

let output = if zk_op::contains_operation(&ZkOperation::Bn128Mul) {
zk_op::ZKVM_OPERATOR.get().unwrap().bn128_run_mul(input).unwrap()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.bn128_run_mul(input)
.unwrap()
} else {
// Not indented to keep the diff clean and make changes to the original code obvious

let input = right_pad::<MUL_INPUT_LEN>(input);
let input = right_pad::<MUL_INPUT_LEN>(input);

let p = read_point(&input[..64])?;
let p = read_point(&input[..64])?;

// `Fr::from_slice` can only fail when the length is not 32.
let fr = bn::Fr::from_slice(&input[64..96]).unwrap();

let mut output = [0u8; 64];
if let Some(mul) = AffineG1::from_jacobian(p * fr) {
mul.x().to_big_endian(&mut output[..32]).unwrap();
mul.y().to_big_endian(&mut output[32..]).unwrap();
}
output
// `Fr::from_slice` can only fail when the length is not 32.
let fr = bn::Fr::from_slice(&input[64..96]).unwrap();

let mut output = [0u8; 64];
if let Some(mul) = AffineG1::from_jacobian(p * fr) {
mul.x().to_big_endian(&mut output[..32]).unwrap();
mul.y().to_big_endian(&mut output[32..]).unwrap();
}
output
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: bn-mul");

Ok((gas_cost, output.into()))
}
Expand All @@ -183,6 +194,8 @@ pub fn run_pair(
pair_base_cost: u64,
gas_limit: u64,
) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: bn-pair");
let gas_used = (input.len() / PAIR_ELEMENT_LEN) as u64 * pair_per_point_cost + pair_base_cost;
if gas_used > gas_limit {
return Err(Error::OutOfGas);
Expand All @@ -193,51 +206,57 @@ pub fn run_pair(
}

let success = if zk_op::contains_operation(&ZkOperation::Bn128Pairing) {
zk_op::ZKVM_OPERATOR.get().unwrap().bn128_run_pairing(input).unwrap()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.bn128_run_pairing(input)
.unwrap()
} else {
// Not indented to keep the diff clean and make changes to the original code obvious

let success = if input.is_empty() {
true
} else {
let elements = input.len() / PAIR_ELEMENT_LEN;

let mut mul = Gt::one();
for idx in 0..elements {
let read_fq_at = |n: usize| {
debug_assert!(n < PAIR_ELEMENT_LEN / 32);
let start = idx * PAIR_ELEMENT_LEN + n * 32;
// SAFETY: We're reading `6 * 32 == PAIR_ELEMENT_LEN` bytes from `input[idx..]`
// per iteration. This is guaranteed to be in-bounds.
let slice = unsafe { input.get_unchecked(start..start + 32) };
Fq::from_slice(slice).map_err(|_| Error::Bn128FieldPointNotAMember)
};
let ax = read_fq_at(0)?;
let ay = read_fq_at(1)?;
let bay = read_fq_at(2)?;
let bax = read_fq_at(3)?;
let bby = read_fq_at(4)?;
let bbx = read_fq_at(5)?;

let a = new_g1_point(ax, ay)?;
let b = {
let ba = Fq2::new(bax, bay);
let bb = Fq2::new(bbx, bby);
if ba.is_zero() && bb.is_zero() {
G2::zero()
} else {
G2::from(AffineG2::new(ba, bb).map_err(|_| Error::Bn128AffineGFailedToCreate)?)
}
};

mul = mul * bn::pairing(a, b);
}

mul == Gt::one()
};

success
let success = if input.is_empty() {
true
} else {
let elements = input.len() / PAIR_ELEMENT_LEN;

let mut mul = Gt::one();
for idx in 0..elements {
let read_fq_at = |n: usize| {
debug_assert!(n < PAIR_ELEMENT_LEN / 32);
let start = idx * PAIR_ELEMENT_LEN + n * 32;
// SAFETY: We're reading `6 * 32 == PAIR_ELEMENT_LEN` bytes from `input[idx..]`
// per iteration. This is guaranteed to be in-bounds.
let slice = unsafe { input.get_unchecked(start..start + 32) };
Fq::from_slice(slice).map_err(|_| Error::Bn128FieldPointNotAMember)
};
let ax = read_fq_at(0)?;
let ay = read_fq_at(1)?;
let bay = read_fq_at(2)?;
let bax = read_fq_at(3)?;
let bby = read_fq_at(4)?;
let bbx = read_fq_at(5)?;

let a = new_g1_point(ax, ay)?;
let b = {
let ba = Fq2::new(bax, bay);
let bb = Fq2::new(bbx, bby);
if ba.is_zero() && bb.is_zero() {
G2::zero()
} else {
G2::from(
AffineG2::new(ba, bb).map_err(|_| Error::Bn128AffineGFailedToCreate)?,
)
}
};

mul = mul * bn::pairing(a, b);
}

mul == Gt::one()
};

success
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: bn-pair");

Ok((gas_used, bool_to_bytes32(success)))
}
Expand Down
40 changes: 28 additions & 12 deletions crates/precompile/src/hash.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,40 +16,56 @@ pub const RIPEMD160: PrecompileWithAddress = PrecompileWithAddress(
/// See: <https://docs.soliditylang.org/en/develop/units-and-global-variables.html#mathematical-and-cryptographic-functions>
/// See: <https://etherscan.io/address/0000000000000000000000000000000000000002>
pub fn sha256_run(input: &Bytes, gas_limit: u64) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: sha256");
let cost = calc_linear_cost_u32(input.len(), 60, 12);
if cost > gas_limit {
let res = if cost > gas_limit {
Err(Error::OutOfGas)
} else {
let output = if zk_op::contains_operation(&ZkOperation::Sha256) {
zk_op::ZKVM_OPERATOR.get().unwrap().sha256_run(input.as_ref()).unwrap().into()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.sha256_run(input.as_ref())
.unwrap()
.into()
} else {
sha2::Sha256::digest(input).to_vec()
};

Ok((cost, output.into()))
}
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: sha256");
res
}

/// See: <https://ethereum.github.io/yellowpaper/paper.pdf>
/// See: <https://docs.soliditylang.org/en/develop/units-and-global-variables.html#mathematical-and-cryptographic-functions>
/// See: <https://etherscan.io/address/0000000000000000000000000000000000000003>
pub fn ripemd160_run(input: &Bytes, gas_limit: u64) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: ripemd160");
let gas_used = calc_linear_cost_u32(input.len(), 600, 120);
if gas_used > gas_limit {
Err(Error::OutOfGas)
} else {
let output = if zk_op::contains_operation(&ZkOperation::Ripemd160) {
zk_op::ZKVM_OPERATOR.get().unwrap().ripemd160_run(input.as_ref()).unwrap().into()
zk_op::ZKVM_OPERATOR
.get()
.unwrap()
.ripemd160_run(input.as_ref())
.unwrap()
.into()
} else {
// Not indented to keep the diff clean and make changes to the original code obvious

let mut hasher = ripemd::Ripemd160::new();
hasher.update(input);
let mut output = [0u8; 32];
hasher.finalize_into((&mut output[12..]).into());
output

let mut hasher = ripemd::Ripemd160::new();
hasher.update(input);
let mut output = [0u8; 32];
hasher.finalize_into((&mut output[12..]).into());
output
};
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: ripemd160");

Ok((gas_used, output.to_vec().into()))
}
Expand Down
4 changes: 4 additions & 0 deletions crates/precompile/src/kzg_point_evaluation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ pub const RETURN_VALUE: &[u8; 64] = &hex!(
/// | 32 | 32 | 32 | 48 | 48 |
/// with z and y being padded 32 byte big endian values
pub fn run(input: &Bytes, gas_limit: u64, env: &Env) -> PrecompileResult {
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-start: kzg");
if gas_limit < GAS_COST {
return Err(Error::OutOfGas);
}
Expand All @@ -49,6 +51,8 @@ pub fn run(input: &Bytes, gas_limit: u64, env: &Env) -> PrecompileResult {
if !verify_kzg_proof(commitment, z, y, proof, env.cfg.kzg_settings.get()) {
return Err(Error::BlobVerifyKzgProofFailed);
}
#[cfg(feature = "sp1-cycle-tracker")]
println!("cycle-tracker-end: kzg");

// Return FIELD_ELEMENTS_PER_BLOB and BLS_MODULUS as padded 32 byte big endian values
Ok((GAS_COST, RETURN_VALUE.into()))
Expand Down
Loading