openvm_circuit_primitives/
hybrid_chip.rs

1use std::marker::PhantomData;
2
3use openvm_cpu_backend::CpuBackend;
4use openvm_cuda_backend::{
5    base::DeviceMatrix, data_transporter::transport_matrix_h2d_col_major,
6    hash_scheme::GpuHashScheme, prelude::SC, GenericGpuBackend, GpuBackend,
7};
8use openvm_cuda_common::stream::GpuDeviceCtx;
9use openvm_stark_backend::prover::{AirProvingContext, ColMajorMatrix};
10
11use crate::Chip;
12
13pub fn get_empty_air_proving_ctx<HS: GpuHashScheme>() -> AirProvingContext<GenericGpuBackend<HS>> {
14    AirProvingContext {
15        cached_mains: vec![],
16        common_main: DeviceMatrix::dummy(),
17        public_values: vec![],
18    }
19}
20
21// Wraps a CPU chip for use with GpuBackend
22pub struct HybridChip<RA, C: Chip<RA, CpuBackend<SC>>> {
23    pub cpu_chip: C,
24    pub device_ctx: GpuDeviceCtx,
25    _marker: PhantomData<RA>,
26}
27
28impl<RA, C: Chip<RA, CpuBackend<SC>>> HybridChip<RA, C> {
29    pub fn new(cpu_chip: C, device_ctx: GpuDeviceCtx) -> Self {
30        Self {
31            cpu_chip,
32            device_ctx,
33            _marker: PhantomData,
34        }
35    }
36}
37
38impl<RA, C: Chip<RA, CpuBackend<SC>>> Chip<RA, GpuBackend> for HybridChip<RA, C> {
39    fn generate_proving_ctx(&self, arena: RA) -> AirProvingContext<GpuBackend> {
40        let ctx = self.cpu_chip.generate_proving_ctx(arena);
41        cpu_proving_ctx_to_gpu(ctx, &self.device_ctx)
42    }
43}
44
45pub fn cpu_proving_ctx_to_gpu<HS: GpuHashScheme>(
46    cpu_ctx: AirProvingContext<CpuBackend<SC>>,
47    device_ctx: &GpuDeviceCtx,
48) -> AirProvingContext<GenericGpuBackend<HS>> {
49    assert!(
50        cpu_ctx.cached_mains.is_empty(),
51        "CPU to GPU transfer of cached traces not supported"
52    );
53    let cm = ColMajorMatrix::from_row_major(&cpu_ctx.common_main);
54    let trace = transport_matrix_h2d_col_major(&cm, device_ctx).unwrap();
55    AirProvingContext {
56        cached_mains: vec![],
57        common_main: trace,
58        public_values: cpu_ctx.public_values,
59    }
60}
openvm_circuit_primitives/hybrid_chip.rs

openvm_circuit_primitives/
hybrid_chip.rs
