Skip to content

Commit

Permalink
Backend artifact (#1706)
Browse files Browse the repository at this point in the history 
Allow the pipeline to setup and cache the backend before calls to
prove/verify.
`Backend` objects are not cloneable, so upon cloning the pipeline, the
instance is cleared (but will be regenerated when needed).
  • Loading branch information
@pacheco
pacheco authored Aug 21, 2024
1 parent 3d4fedf commit 3f3b291
Showing 1 changed file with 69 additions and 137 deletions.
206 changes: 69 additions & 137 deletions pipeline/src/pipeline.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use powdr_ast::{
object::PILGraph,
parsed::{asm::ASMProgram, PILFile},
};
use powdr_backend::{BackendOptions, BackendType, Proof};
use powdr_backend::{Backend, BackendOptions, BackendType, Proof};
use powdr_executor::{
constant_evaluator::{self, get_uniquely_sized_cloned, VariablySizedColumn},
witgen::{
Expand All @@ -39,7 +39,7 @@ use std::collections::BTreeMap;
type Columns<T> = Vec<(String, Vec<T>)>;
type VariablySizedColumns<T> = Vec<(String, VariablySizedColumn<T>)>;

#[derive(Default, Clone)]
#[derive(Default)]
pub struct Artifacts<T: FieldElement> {
/// The path to a single .asm file.
asm_file_path: Option<PathBuf>,
Expand Down Expand Up @@ -72,6 +72,8 @@ pub struct Artifacts<T: FieldElement> {
fixed_cols: Option<Arc<VariablySizedColumns<T>>>,
/// Generated witnesses.
witness: Option<Arc<Columns<T>>>,
/// Instantiated backend.
backend: Option<Box<dyn Backend<T>>>,
/// The proof (if successful).
proof: Option<Proof>,
}
Expand Down Expand Up @@ -138,6 +140,30 @@ pub struct Pipeline<T: FieldElement> {
host_context: HostContext,
}

impl<T: FieldElement> Clone for Artifacts<T> {
fn clone(&self) -> Self {
Artifacts {
asm_file_path: self.asm_file_path.clone(),
asm_string: self.asm_string.clone(),
parsed_asm_file: self.parsed_asm_file.clone(),
resolved_module_tree: self.resolved_module_tree.clone(),
analyzed_asm: self.analyzed_asm.clone(),
constrained_machine_collection: self.constrained_machine_collection.clone(),
linked_machine_graph: self.linked_machine_graph.clone(),
parsed_pil_file: self.parsed_pil_file.clone(),
pil_file_path: self.pil_file_path.clone(),
pil_string: self.pil_string.clone(),
analyzed_pil: self.analyzed_pil.clone(),
optimized_pil: self.optimized_pil.clone(),
fixed_cols: self.fixed_cols.clone(),
witness: self.witness.clone(),
proof: self.proof.clone(),
// Backend is not cloneable, so we clear it instead
backend: None,
}
}
}

use super::HostContext;

impl<T> Default for Pipeline<T>
Expand Down Expand Up @@ -280,21 +306,25 @@ impl<T: FieldElement> Pipeline<T> {
pub fn with_backend(mut self, backend: BackendType, options: Option<BackendOptions>) -> Self {
self.arguments.backend = Some(backend);
self.arguments.backend_options = options.unwrap_or_default();
self.artifact.backend = None;
self
}

pub fn with_setup_file(mut self, setup_file: Option<PathBuf>) -> Self {
self.arguments.setup_file = setup_file;
self.artifact.backend = None;
self
}

pub fn with_vkey_file(mut self, vkey_file: Option<PathBuf>) -> Self {
self.arguments.vkey_file = vkey_file;
self.artifact.backend = None;
self
}

pub fn with_vkey_app_file(mut self, vkey_app_file: Option<PathBuf>) -> Self {
self.arguments.vkey_app_file = vkey_app_file;
self.artifact.backend = None;
self
}

Expand Down Expand Up @@ -885,20 +915,18 @@ impl<T: FieldElement> Pipeline<T> {
)))
}

pub fn compute_proof(&mut self) -> Result<&Proof, Vec<String>> {
if self.artifact.proof.is_some() {
return Ok(self.artifact.proof.as_ref().unwrap());
}
pub fn backend(&mut self) -> Result<&mut dyn Backend<T>, Vec<String>> {
Ok(self.artifact.backend.as_deref_mut().unwrap())
}

pub fn setup_backend(&mut self) -> Result<&mut dyn Backend<T>, Vec<String>> {
if self.artifact.backend.is_some() {
return Ok(self.artifact.backend.as_deref_mut().unwrap());
}
let pil = self.compute_optimized_pil()?;
let fixed_cols = self.compute_fixed_cols()?;
let witness = self.compute_witness()?;
let witgen_callback = self.witgen_callback()?;

let backend = self
.arguments
.backend
.expect("backend must be set before calling proving!");
let backend = self.arguments.backend.expect("no backend selected!");
let factory = backend.factory::<T>();

// Opens the setup file, if set.
Expand All @@ -922,7 +950,9 @@ impl<T: FieldElement> Pipeline<T> {
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

/* Create the backend */
// Create the backend
let start = Instant::now();
self.log(&format!("Backend setup for {backend}..."));
let backend = factory
.create(
pil.clone(),
Expand All @@ -934,6 +964,19 @@ impl<T: FieldElement> Pipeline<T> {
self.arguments.backend_options.clone(),
)
.unwrap();
self.log(&format!("Setup took {}s", start.elapsed().as_secs_f32()));

self.artifact.backend = Some(backend);
Ok(self.artifact.backend.as_deref_mut().unwrap())
}

pub fn compute_proof(&mut self) -> Result<&Proof, Vec<String>> {
if self.artifact.proof.is_some() {
return Ok(self.artifact.proof.as_ref().unwrap());
}

let witness = self.compute_witness()?;
let witgen_callback = self.witgen_callback()?;

// Reads the existing proof file, if set.
let existing_proof = self
Expand All @@ -942,22 +985,25 @@ impl<T: FieldElement> Pipeline<T> {
.as_ref()
.map(|path| fs::read(path).unwrap());

self.setup_backend()?;

let start = Instant::now();
let proof = match backend.prove(&witness, existing_proof, witgen_callback) {
Ok(proof) => proof,
Err(powdr_backend::Error::BackendError(e)) => {
return Err(vec![e.to_string()]);
let proof = {
let backend = self.backend()?;
match backend.prove(&witness, existing_proof, witgen_callback) {
Ok(proof) => proof,
Err(powdr_backend::Error::BackendError(e)) => {
return Err(vec![e.to_string()]);
}
Err(e) => panic!("{}", e),
}
Err(e) => panic!("{}", e),
};
self.log(&format!(
"Proof generation took {}s",
start.elapsed().as_secs_f32()
));
self.log(&format!("Proof size: {} bytes", proof.len()));

drop(backend);

self.maybe_write_proof(&proof)?;

self.artifact.proof = Some(proof);
Expand Down Expand Up @@ -989,44 +1035,7 @@ impl<T: FieldElement> Pipeline<T> {
&mut self,
mut writer: W,
) -> Result<(), Vec<String>> {
let backend = self
.arguments
.backend
.expect("backend must be set before generating verification key!");
let factory = backend.factory::<T>();

let mut setup_file = self
.arguments
.setup_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

// An aggregation verification key needs the app vkey to be set
let mut vkey_app_file = self
.arguments
.vkey_app_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

let pil = self.compute_optimized_pil()?;
let fixed_cols = self.compute_fixed_cols()?;

let backend = factory
.create(
pil.clone(),
fixed_cols.clone(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
None,
vkey_app_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
self.arguments.backend_options.clone(),
)
.unwrap();

let backend = self.setup_backend()?;
match backend.export_verification_key(&mut writer) {
Ok(()) => Ok(()),
Err(powdr_backend::Error::BackendError(e)) => Err(vec![e]),
Expand All @@ -1038,46 +1047,7 @@ impl<T: FieldElement> Pipeline<T> {
&mut self,
mut writer: W,
) -> Result<(), Vec<String>> {
let backend = self
.arguments
.backend
.expect("backend must be set before generating verifier!");
let factory = backend.factory::<T>();

let mut setup_file = self
.arguments
.setup_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

let mut vkey = self
.arguments
.vkey_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

let mut vkey_app = self
.arguments
.vkey_app_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

let pil = self.compute_optimized_pil()?;
let fixed_cols = self.compute_fixed_cols()?;

let backend = factory
.create(
pil.clone(),
fixed_cols.clone(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
vkey.as_io_read(),
vkey_app.as_io_read(),
self.arguments.backend_options.clone(),
)
.unwrap();
let backend = self.setup_backend()?;

match backend.export_ethereum_verifier(&mut writer) {
Ok(()) => Ok(()),
Expand All @@ -1090,45 +1060,7 @@ impl<T: FieldElement> Pipeline<T> {
}

pub fn verify(&mut self, proof: &[u8], instances: &[Vec<T>]) -> Result<(), Vec<String>> {
let backend = self
.arguments
.backend
.expect("backend must be set before generating verification key!");
let factory = backend.factory::<T>();

let mut setup_file = self
.arguments
.setup_file
.as_ref()
.map(|path| BufReader::new(fs::File::open(path).unwrap()));

let mut vkey_file = self
.arguments
.vkey_file
.as_ref()
.map(fs::File::open)
.map(Result::unwrap)
.map(BufReader::new);

let pil = self.compute_optimized_pil()?;
let fixed_cols = self.compute_fixed_cols()?;

let backend = factory
.create(
pil.clone(),
fixed_cols.clone(),
self.output_dir.clone(),
setup_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
vkey_file
.as_mut()
.map(|file| file as &mut dyn std::io::Read),
// We shouldn't need the app verification key for this
None,
self.arguments.backend_options.clone(),
)
.unwrap();
let backend = self.setup_backend()?;

let start = Instant::now();
match backend.verify(proof, instances) {
Expand Down

0 comments on commit 3f3b291

Please sign in to comment.