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}