openvm_stark_backend/prover/
hal.rs

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