openvm_stark_backend/prover/hal.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
//! # Hardware Abstraction Layer
//!
//! Not all hardware implementations need to implement this.
//! A pure external device implementation can just implement the [Prover](super::Prover) trait directly.
use std::sync::Arc;
use p3_challenger::CanObserve;
use p3_matrix::dense::RowMajorMatrix;
use serde::{de::DeserializeOwned, Serialize};
use super::types::{
DeviceMultiStarkProvingKey, DeviceStarkProvingKey, PairView, ProverDataAfterRapPhases,
SingleCommitPreimage,
};
use crate::{
config::{Com, StarkGenericConfig, Val},
keygen::types::MultiStarkProvingKey,
};
/// Associated types needed by the prover, in the form of buffers and views,
/// specific to a specific hardware backend.
///
/// Memory allocation and copying is not handled by this trait.
pub trait ProverBackend {
/// Extension field degree for the challenge field `Self::Challenge` over base field `Self::Val`.
const CHALLENGE_EXT_DEGREE: u8;
// ==== Host Types ====
/// Base field type, on host.
type Val: Copy + Send + Sync + Serialize + DeserializeOwned;
/// Challenge field (extension field of base field), on host.
type Challenge: Copy + Send + Sync + Serialize + DeserializeOwned;
/// PCS opening proof on host (see [OpeningProver]). This should not be a reference.
type OpeningProof: Clone + Send + Sync + Serialize + DeserializeOwned;
/// Partial proof for multiple RAPs
type RapPartialProof: Clone + Send + Sync + Serialize + DeserializeOwned;
/// Single commitment on host.
// Commitments are small in size and need to be transferred back to host to be included in proof.
type Commitment: Clone + Send + Sync + Serialize + DeserializeOwned;
/// Challenger to observe commitments. Sampling is left to other trait implementations.
/// We anticipate that the challenger largely operates on the host.
type Challenger: CanObserve<Self::Val> + CanObserve<Self::Commitment>;
// ==== Device Types ====
/// Single matrix buffer on device together with dimension metadata. Owning this means nothing else has a shared
/// reference to the buffer.
type Matrix: MatrixDimensions + Send + Sync;
/// Owned buffer for the preimage of a PCS commitment on device, together with any metadata
/// necessary for computing opening proofs.
///
/// For example, multiple buffers for LDE matrices, their trace domain sizes, and pointer to mixed merkle tree.
type PcsData: Send + Sync;
/// Part of proving key for a single RAP specific for the RAP challenge phases
type RapPartialProvingKey: Send + Sync;
}
pub trait MatrixDimensions {
fn height(&self) -> usize;
fn width(&self) -> usize;
}
pub trait ProverDevice<PB: ProverBackend>:
TraceCommitter<PB> + RapPartialProver<PB> + QuotientCommitter<PB> + OpeningProver<PB>
{
}
/// Provides functionality for committing to a batch of trace matrices, possibly of different heights.
pub trait TraceCommitter<PB: ProverBackend> {
fn commit(&self, traces: &[PB::Matrix]) -> (PB::Commitment, PB::PcsData);
}
/// This trait is responsible for all partial proving of after challenge rounds (a.k.a layers) in a
/// RAP after the main trace has been committed.
///
/// The partial prover *may*:
/// - observe and/or sample challenges
/// - commit to additional trace data
/// - generate other partial proof data
pub trait RapPartialProver<PB: ProverBackend> {
/// The `trace_views` are the views of the respective trace matrices, evaluated on the trace domain.
/// Currently this function does not provide a view of any already committed data associated
/// with the trace views, although that data is available.
fn partially_prove<'a>(
&self,
challenger: &mut PB::Challenger,
mpk: &DeviceMultiStarkProvingKey<'a, PB>,
trace_views: Vec<PairView<&'a PB::Matrix, PB::Val>>,
) -> (PB::RapPartialProof, ProverDataAfterRapPhases<PB>);
}
/// Only needed in proof systems that use quotient polynomials.
pub trait QuotientCommitter<PB: ProverBackend> {
/// Given a view of the PCS data from all phases of proving,
/// first get the trace polynomials evaluated on the quotient domains.
/// Then compute the quotient polynomial evaluated on the quotient domain
/// and commit to it.
///
/// The lengths of
/// - `pk_views`: proving key per AIR
/// - `public_values`: public values per AIR
/// - `cached_views_per_air`: committed trace views per AIR (if any)
///
/// must be equal, and all equal to the number of AIRs.
///
/// Quotient polynomials for multiple RAP matrices are committed together into a single commitment.
/// The quotient polynomials can be committed together even if the corresponding trace matrices
/// are committed separately.
fn eval_and_commit_quotient(
&self,
challenger: &mut PB::Challenger,
pk_views: &[DeviceStarkProvingKey<PB>],
public_values: &[Vec<PB::Val>],
cached_views_per_air: &[Vec<SingleCommitPreimage<&PB::Matrix, &PB::PcsData>>],
common_main_pcs_data: &PB::PcsData,
prover_data_after: &ProverDataAfterRapPhases<PB>,
) -> (PB::Commitment, PB::PcsData);
}
/// Polynomial commitment scheme (PCS) opening proof generator.
pub trait OpeningProver<PB: ProverBackend> {
/// Opening proof for multiple RAP matrices, where
/// - (for now) each preprocessed trace matrix has a separate commitment
/// - main trace matrices can have multiple commitments
/// - for each after_challenge phase, all matrices in the phase share a commitment
/// - quotient poly chunks are all committed together
// Note[jpw]: pass `preprocessed, main` by reference because there is cached data
// that is not owned. We assume `after_phase, quotient_data` will never be used after this.
fn open(
&self,
challenger: &mut PB::Challenger,
// For each preprocessed trace commitment, the prover data and
// the log height of the matrix, in order
preprocessed: Vec<&PB::PcsData>,
// For each main trace commitment, the prover data and
// the log height of each matrix, in order
// Note: this is all one challenge phase.
main: Vec<&PB::PcsData>,
// `after_phase[i]` has shared commitment prover data for all matrices in phase `i + 1`.
after_phase: Vec<PB::PcsData>,
// Quotient poly commitment prover data
quotient_data: PB::PcsData,
// Quotient degree for each RAP committed in quotient_data, in order
quotient_degrees: &[u8],
) -> PB::OpeningProof;
}
/// Trait to manage data transport of prover types from host to device.
pub trait DeviceDataTransporter<SC, PB>
where
SC: StarkGenericConfig,
PB: ProverBackend<Val = Val<SC>, Challenge = SC::Challenge, Commitment = Com<SC>>,
{
/// Transport the proving key to the device, filtering for only the provided `air_ids`.
fn transport_pk_to_device<'a>(
&self,
mpk: &'a MultiStarkProvingKey<SC>,
air_ids: Vec<usize>,
) -> DeviceMultiStarkProvingKey<'a, PB>
where
SC: 'a;
fn transport_matrix_to_device(&self, matrix: &Arc<RowMajorMatrix<Val<SC>>>) -> PB::Matrix;
fn transport_pcs_data_to_device(&self, data: &super::cpu::PcsData<SC>) -> PB::PcsData;
}