openvm_cuda_backend/
data_transporter.rs

1use std::{cmp::max, fmt::Debug, sync::Arc};
2
3use itertools::Itertools;
4use openvm_cuda_common::{
5    copy::{cuda_memcpy_on, MemCopyD2H, MemCopyH2D},
6    d_buffer::DeviceBuffer,
7    error::MemCopyError,
8    stream::{sync_stream, GpuDeviceCtx},
9};
10use openvm_stark_backend::{
11    keygen::types::MultiStarkProvingKey,
12    p3_matrix::{dense::RowMajorMatrix, Matrix},
13    prover::{
14        stacked_pcs::{MerkleTree, StackedPcsData},
15        AirProvingContext, ColMajorMatrix, CommittedTraceData, CpuColMajorBackend,
16        DeviceDataTransporter, DeviceMultiStarkProvingKey, DeviceStarkProvingKey, MatrixDimensions,
17        MatrixView, ProvingContext,
18    },
19};
20use tracing::{debug, info_span};
21
22use crate::{
23    base::DeviceMatrix,
24    cuda::matrix::{collapse_strided_matrix, matrix_transpose_fp},
25    gpu_backend::GenericGpuBackend,
26    hash_scheme::GpuHashScheme,
27    merkle_tree::MerkleTreeGpu,
28    poly::PleMatrix,
29    prelude::{Digest, F, SC},
30    stacked_pcs::StackedPcsDataGpu,
31    AirDataGpu, GpuBackend, GpuDevice, GpuProverConfig, ProverError,
32};
33
34impl<HS: GpuHashScheme> DeviceDataTransporter<HS::SC, GenericGpuBackend<HS>> for GpuDevice {
35    fn transport_pk_to_device(
36        &self,
37        mpk: &MultiStarkProvingKey<HS::SC>,
38    ) -> DeviceMultiStarkProvingKey<GenericGpuBackend<HS>> {
39        let device_ctx = &self.device_ctx;
40        let _span = info_span!("transport_pk_to_device").entered();
41        let per_air = mpk
42            .per_air
43            .iter()
44            .map(|pk| {
45                let preprocessed_data = pk.preprocessed_data.as_ref().map(|d| {
46                    transport_and_unstack_single_data_h2d::<HS>(
47                        d.as_ref(),
48                        self.prover_config(),
49                        device_ctx,
50                    )
51                    .unwrap()
52                });
53                let other_data = AirDataGpu::new(pk, device_ctx).unwrap();
54                let num_monomials = other_data
55                    .zerocheck_monomials
56                    .as_ref()
57                    .map(|m| m.num_monomials)
58                    .unwrap_or(0);
59                debug!(air = %pk.air_name, num_monomials, "monomial expansion");
60
61                DeviceStarkProvingKey {
62                    air_name: pk.air_name.clone(),
63                    vk: pk.vk.clone(),
64                    preprocessed_data,
65                    other_data,
66                }
67            })
68            .collect();
69        // Synchronize in case the proving key is shared between threads/streams
70        self.device_ctx.stream.synchronize().unwrap();
71
72        DeviceMultiStarkProvingKey::new(
73            per_air,
74            mpk.trace_height_constraints.clone(),
75            mpk.max_constraint_degree,
76            mpk.params.clone(),
77            mpk.vk_pre_hash,
78        )
79    }
80
81    fn transport_matrix_to_device(&self, matrix: &ColMajorMatrix<F>) -> DeviceMatrix<F> {
82        transport_matrix_h2d_col_major(matrix, &self.device_ctx).unwrap()
83    }
84
85    fn transport_pcs_data_to_device(
86        &self,
87        pcs_data: &StackedPcsData<F, HS::Digest>,
88    ) -> StackedPcsDataGpu<F, HS::Digest> {
89        transport_pcs_data_h2d::<HS::Digest>(pcs_data, self.prover_config(), &self.device_ctx)
90            .unwrap()
91    }
92
93    fn transport_matrix_from_device_to_host(&self, matrix: &DeviceMatrix<F>) -> ColMajorMatrix<F> {
94        transport_matrix_d2h_col_major(matrix, &self.device_ctx).unwrap()
95    }
96}
97
98pub fn transport_matrix_h2d_col_major<T>(
99    matrix: &ColMajorMatrix<T>,
100    device_ctx: &GpuDeviceCtx,
101) -> Result<DeviceMatrix<T>, MemCopyError> {
102    // matrix is already col-major, so this is just H2D buffer transfer
103    let buffer = matrix.values.to_device_on(device_ctx)?;
104    unsafe { sync_stream(device_ctx.stream.as_raw())? };
105    Ok(DeviceMatrix::new(
106        Arc::new(buffer),
107        matrix.height(),
108        matrix.width(),
109    ))
110}
111
112pub fn transport_matrix_h2d_row(
113    matrix: &RowMajorMatrix<F>,
114    device_ctx: &GpuDeviceCtx,
115) -> Result<DeviceMatrix<F>, MemCopyError> {
116    let data = matrix.values.as_slice();
117    let input_buffer = data.to_device_on(device_ctx).unwrap();
118    let output = DeviceMatrix::<F>::with_capacity_on(
119        Matrix::height(matrix),
120        Matrix::width(matrix),
121        device_ctx,
122    );
123    unsafe {
124        matrix_transpose_fp(
125            output.buffer(),
126            &input_buffer,
127            Matrix::width(matrix),
128            Matrix::height(matrix),
129            device_ctx.stream.as_raw(),
130        )?;
131    }
132    unsafe { sync_stream(device_ctx.stream.as_raw())? };
133    assert_eq!(output.strong_count(), 1);
134    Ok(output)
135}
136
137/// `d` must be the stacked pcs data of a single trace matrix.
138/// This function will transport `d` to device and then unstack it (allocating device memory) to
139/// return `CommittedTraceData<F, Digest>`.
140pub fn transport_and_unstack_single_data_h2d<HS: GpuHashScheme>(
141    d: &StackedPcsData<F, HS::Digest>,
142    prover_config: &GpuProverConfig,
143    device_ctx: &GpuDeviceCtx,
144) -> Result<CommittedTraceData<GenericGpuBackend<HS>>, ProverError> {
145    let _span = info_span!("transport_unstack_h2d").entered();
146    debug_assert!(d
147        .layout
148        .sorted_cols
149        .iter()
150        .all(|(mat_idx, _, _)| *mat_idx == 0));
151    let l_skip = d.layout.l_skip();
152    let trace_view = d.mat_view(0);
153    let height = trace_view.height();
154    let width = trace_view.width();
155    let stride = trace_view.stride();
156    let lifted_height = height * stride;
157    debug_assert_eq!(lifted_height, max(height, 1 << l_skip));
158    debug_assert_eq!(lifted_height * width, trace_view.values().len());
159    debug_assert!(d.matrix.values.len() >= lifted_height * width);
160    let stacked_width = d.matrix.width();
161    let stacked_height = d.matrix.height();
162    let d_matrix_evals = d.matrix.values.to_device_on(device_ctx)?;
163    let strided_trace = DeviceBuffer::<F>::with_capacity_on(lifted_height * width, device_ctx);
164    // SAFETY: D2D copy
165    // - `d_matrix_evals` is the stacked matrix, guaranteed to have length `>= lifted_height *
166    //   width` by definition of stacking
167    // - `d_matrix_evals` stacks a single trace matrix
168    unsafe {
169        cuda_memcpy_on::<true, true>(
170            strided_trace.as_mut_raw_ptr(),
171            d_matrix_evals.as_raw_ptr(),
172            lifted_height * width * size_of::<F>(),
173            device_ctx,
174        )?;
175    }
176    let trace_buffer = if stride == 1 {
177        strided_trace
178    } else {
179        let buf = DeviceBuffer::<F>::with_capacity_on(height * width, device_ctx);
180        unsafe {
181            collapse_strided_matrix(
182                buf.as_mut_ptr(),
183                strided_trace.as_ptr(),
184                width as u32,
185                height as u32,
186                stride as u32,
187                device_ctx.stream.as_raw(),
188            )
189            .map_err(ProverError::CollapseStrided)?;
190        }
191        // Wait for kernel to finish so we can safely drop strided_trace
192        unsafe { sync_stream(device_ctx.stream.as_raw()) }
193            .map_err(ProverError::StreamSynchronize)?;
194        drop(strided_trace);
195        buf
196    };
197    let d_matrix = prover_config.cache_stacked_matrix.then(|| {
198        PleMatrix::from_evals(
199            l_skip,
200            d_matrix_evals,
201            stacked_height,
202            stacked_width,
203            device_ctx,
204        )
205    });
206    let d_tree =
207        transport_merkle_tree_h2d(&d.tree, prover_config.cache_rs_code_matrix, device_ctx)?;
208    let d_data = StackedPcsDataGpu {
209        layout: d.layout.clone(),
210        matrix: d_matrix,
211        tree: d_tree,
212    };
213    // Sanity check. Not a strong assert because we transport the merkle tree
214    // instead of recomputing it above.
215    assert!(
216        d_data.tree.root() == d.commit().unwrap(),
217        "transported tree root mismatch"
218    );
219    Ok(CommittedTraceData {
220        commitment: d.commit().unwrap(),
221        trace: DeviceMatrix::new(Arc::new(trace_buffer), height, width),
222        data: Arc::new(d_data),
223    })
224}
225
226/// Transports backing matrix and tree digest layers from host to device.
227pub fn transport_merkle_tree_h2d<F, Digest: Clone>(
228    tree: &MerkleTree<F, Digest>,
229    cache_backing_matrix: bool,
230    device_ctx: &GpuDeviceCtx,
231) -> Result<MerkleTreeGpu<F, Digest>, MemCopyError> {
232    let backing_matrix = if cache_backing_matrix {
233        Some(transport_matrix_h2d_col_major(
234            tree.backing_matrix(),
235            device_ctx,
236        )?)
237    } else {
238        None
239    };
240    let digest_layers = tree
241        .digest_layers()
242        .iter()
243        .map(|layer| layer.to_device_on(device_ctx))
244        .collect::<Result<Vec<_>, _>>()?;
245    unsafe { sync_stream(device_ctx.stream.as_raw())? };
246    Ok(MerkleTreeGpu {
247        backing_matrix,
248        digest_layers,
249        rows_per_query: tree.rows_per_query(),
250        root: tree.root().unwrap(),
251    })
252}
253
254pub fn transport_pcs_data_h2d<D: Copy + Clone + PartialEq + Send + Sync + 'static>(
255    pcs_data: &StackedPcsData<F, D>,
256    prover_config: &GpuProverConfig,
257    device_ctx: &GpuDeviceCtx,
258) -> Result<StackedPcsDataGpu<F, D>, ProverError> {
259    let _span = info_span!("transport_pcs_data_h2d").entered();
260    let StackedPcsData {
261        layout,
262        matrix,
263        tree,
264    } = pcs_data;
265    let width = matrix.width();
266    let height = matrix.height();
267    let d_matrix_evals = matrix.values.to_device_on(device_ctx)?;
268    let d_matrix = prover_config
269        .cache_stacked_matrix
270        .then(|| PleMatrix::from_evals(layout.l_skip(), d_matrix_evals, height, width, device_ctx));
271    let d_tree = transport_merkle_tree_h2d(tree, prover_config.cache_rs_code_matrix, device_ctx)?;
272    unsafe { sync_stream(device_ctx.stream.as_raw()) }.map_err(ProverError::StreamSynchronize)?;
273
274    Ok(StackedPcsDataGpu {
275        layout: layout.clone(),
276        matrix: d_matrix,
277        tree: d_tree,
278    })
279}
280
281pub fn transport_air_proving_ctx_to_device<HS: GpuHashScheme>(
282    cpu_ctx: AirProvingContext<CpuColMajorBackend<SC>>,
283    device_ctx: &GpuDeviceCtx,
284) -> AirProvingContext<GenericGpuBackend<HS>> {
285    let _span = info_span!("transport_air_ctx_h2d").entered();
286    assert!(
287        cpu_ctx.cached_mains.is_empty(),
288        "CPU to GPU transfer of cached traces not supported"
289    );
290    let trace = transport_matrix_h2d_col_major(&cpu_ctx.common_main, device_ctx).unwrap();
291    AirProvingContext {
292        cached_mains: vec![],
293        common_main: trace,
294        public_values: cpu_ctx.public_values,
295    }
296}
297
298pub fn transport_proving_ctx_to_host(
299    gpu_ctx: ProvingContext<GpuBackend>,
300    l_skip: usize,
301    device_ctx: &GpuDeviceCtx,
302) -> ProvingContext<CpuColMajorBackend<SC>> {
303    let per_trace = gpu_ctx
304        .per_trace
305        .into_iter()
306        .map(|(i, air_ctx)| {
307            (
308                i,
309                transport_air_proving_ctx_to_host(air_ctx, l_skip, device_ctx),
310            )
311        })
312        .collect_vec();
313    ProvingContext { per_trace }
314}
315
316pub fn transport_air_proving_ctx_to_host(
317    gpu_ctx: AirProvingContext<GpuBackend>,
318    l_skip: usize,
319    device_ctx: &GpuDeviceCtx,
320) -> AirProvingContext<CpuColMajorBackend<SC>> {
321    let trace = transport_matrix_d2h_col_major(&gpu_ctx.common_main, device_ctx).unwrap();
322    let cached_mains = gpu_ctx
323        .cached_mains
324        .into_iter()
325        .map(|mat| {
326            // WARNING: By default this matrix isn't cached. For this to work, ensure that
327            // GpuProverConfig fields cache_stacked_matrix and cache_rs_code_matrix are set
328            // to be true.
329            let evals_matrix = mat
330                .data
331                .matrix
332                .as_ref()
333                .unwrap()
334                .to_evals(l_skip, device_ctx)
335                .unwrap();
336            CommittedTraceData {
337                commitment: mat.commitment,
338                trace: transport_matrix_d2h_col_major(&mat.trace, device_ctx).unwrap(),
339                data: Arc::new(StackedPcsData {
340                    layout: mat.data.layout.clone(),
341                    matrix: transport_matrix_d2h_col_major(&evals_matrix, device_ctx).unwrap(),
342                    tree: transport_merkle_tree_to_host(&mat.data.tree, device_ctx),
343                }),
344            }
345        })
346        .collect_vec();
347    AirProvingContext {
348        cached_mains,
349        common_main: trace,
350        public_values: gpu_ctx.public_values,
351    }
352}
353
354pub fn transport_matrix_d2h_col_major<T>(
355    matrix: &DeviceMatrix<T>,
356    device_ctx: &GpuDeviceCtx,
357) -> Result<ColMajorMatrix<T>, MemCopyError> {
358    let values_host = matrix.buffer().to_host_on(device_ctx)?;
359    Ok(ColMajorMatrix::new(values_host, matrix.width()))
360}
361
362pub fn transport_matrix_d2h_row_major(
363    matrix: &DeviceMatrix<F>,
364    device_ctx: &GpuDeviceCtx,
365) -> Result<RowMajorMatrix<F>, MemCopyError> {
366    let matrix_buffer =
367        DeviceBuffer::<F>::with_capacity_on(matrix.height() * matrix.width(), device_ctx);
368    unsafe {
369        matrix_transpose_fp(
370            &matrix_buffer,
371            matrix.buffer(),
372            matrix.height(),
373            matrix.width(),
374            device_ctx.stream.as_raw(),
375        )?;
376    }
377    Ok(RowMajorMatrix::<F>::new(
378        matrix_buffer.to_host_on(device_ctx)?,
379        matrix.width(),
380    ))
381}
382
383/// For debugging purposes only.
384///
385/// # Panics
386/// If `tree.backing_matrix` is `None`.
387pub fn transport_merkle_tree_to_host(
388    tree: &MerkleTreeGpu<F, Digest>,
389    device_ctx: &GpuDeviceCtx,
390) -> MerkleTree<F, Digest> {
391    let backing_matrix =
392        transport_matrix_d2h_col_major(tree.backing_matrix.as_ref().unwrap(), device_ctx).unwrap();
393    let digest_layers = tree
394        .digest_layers
395        .iter()
396        .map(|layer| layer.to_host_on(device_ctx).unwrap())
397        .collect_vec();
398    // Safety: assuming the tree is properly constructed on device, the layers are correct after D2H
399    // transfer.
400    unsafe {
401        MerkleTree::<F, Digest>::from_raw_parts(backing_matrix, digest_layers, tree.rows_per_query)
402    }
403}
404
405pub fn assert_eq_host_and_device_matrix_col_maj<T: Clone + Send + Sync + PartialEq + Debug>(
406    cpu: &ColMajorMatrix<T>,
407    gpu: &DeviceMatrix<T>,
408    device_ctx: &GpuDeviceCtx,
409) {
410    assert_eq!(gpu.width(), cpu.width());
411    assert_eq!(gpu.height(), cpu.height());
412    let gpu = gpu.buffer().to_host_on(device_ctx).unwrap();
413    for r in 0..cpu.height() {
414        for c in 0..cpu.width() {
415            assert_eq!(
416                gpu[c * cpu.height() + r],
417                *cpu.get(r, c).unwrap(),
418                "Mismatch at row {} column {}",
419                r,
420                c
421            );
422        }
423    }
424}
425
426pub fn assert_eq_device_matrix<T: Clone + Send + Sync + PartialEq + Debug>(
427    a: &DeviceMatrix<T>,
428    b: &DeviceMatrix<T>,
429    device_ctx: &GpuDeviceCtx,
430) {
431    assert_eq!(a.height(), b.height());
432    assert_eq!(a.width(), b.width());
433    assert_eq!(a.buffer().len(), b.buffer().len());
434    let a_host = a.buffer().to_host_on(device_ctx).unwrap();
435    let b_host = b.buffer().to_host_on(device_ctx).unwrap();
436    for r in 0..a.height() {
437        for c in 0..a.width() {
438            assert_eq!(
439                a_host[c * a.height() + r],
440                b_host[c * b.height() + r],
441                "Mismatch at row {} column {}",
442                r,
443                c
444            );
445        }
446    }
447}
448
449pub fn assert_eq_host_and_device_matrix<T: Clone + Send + Sync + PartialEq + Debug>(
450    cpu: Arc<RowMajorMatrix<T>>,
451    gpu: &DeviceMatrix<T>,
452    device_ctx: &GpuDeviceCtx,
453) {
454    assert_eq!(gpu.width(), Matrix::width(cpu.as_ref()));
455    assert_eq!(gpu.height(), Matrix::height(cpu.as_ref()));
456    let gpu = gpu.buffer().to_host_on(device_ctx).unwrap();
457    for r in 0..Matrix::height(cpu.as_ref()) {
458        for c in 0..Matrix::width(cpu.as_ref()) {
459            assert_eq!(
460                gpu[c * Matrix::height(cpu.as_ref()) + r],
461                cpu.get(r, c).expect("matrix index out of bounds"),
462                "Mismatch at row {} column {}",
463                r,
464                c
465            );
466        }
467    }
468}
openvm_cuda_backend/data_transporter.rs

openvm_cuda_backend/
data_transporter.rs
