openvm_circuit/system/cuda/
program.rs

1use std::{mem::size_of, sync::Arc};
2
3use openvm_circuit::{primitives::Chip, system::program::ProgramExecutionCols};
4use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend, GpuDevice};
5use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
6use openvm_instructions::{
7    program::{Program, DEFAULT_PC_STEP},
8    LocalOpcode, SystemOpcode,
9};
10use openvm_stark_backend::prover::{
11    AirProvingContext, CommittedTraceData, MatrixDimensions, TraceCommitter,
12};
13use p3_field::PrimeCharacteristicRing;
14
15use crate::cuda_abi::program;
16
17pub struct ProgramChipGPU {
18    pub cached: Option<CommittedTraceData<GpuBackend>>,
19    pub device_ctx: GpuDeviceCtx,
20}
21
22impl ProgramChipGPU {
23    pub fn new(device_ctx: GpuDeviceCtx) -> Self {
24        Self {
25            cached: None,
26            device_ctx,
27        }
28    }
29
30    pub fn generate_cached_trace(
31        program: Program<F>,
32        device_ctx: &GpuDeviceCtx,
33    ) -> DeviceMatrix<F> {
34        let instructions = program
35            .enumerate_by_pc()
36            .into_iter()
37            .map(|(pc, instruction, _)| {
38                [
39                    F::from_u32(pc),
40                    instruction.opcode.to_field(),
41                    instruction.a,
42                    instruction.b,
43                    instruction.c,
44                    instruction.d,
45                    instruction.e,
46                    instruction.f,
47                    instruction.g,
48                ]
49            })
50            .collect::<Vec<_>>();
51
52        let num_records = instructions.len();
53        let height = num_records.next_power_of_two();
54        let records = instructions
55            .into_iter()
56            .flatten()
57            .collect::<Vec<_>>()
58            .to_device_on(device_ctx)
59            .unwrap();
60
61        let trace = DeviceMatrix::<F>::with_capacity_on(
62            height,
63            size_of::<ProgramExecutionCols<u8>>(),
64            device_ctx,
65        );
66        trace.buffer().fill_zero_on(device_ctx).unwrap();
67        unsafe {
68            program::cached_tracegen(
69                trace.buffer(),
70                trace.height(),
71                trace.width(),
72                &records,
73                program.pc_base,
74                DEFAULT_PC_STEP,
75                SystemOpcode::TERMINATE.global_opcode().as_usize(),
76                device_ctx.stream.as_raw(),
77            )
78            .expect("Failed to generate cached trace");
79        }
80        trace
81    }
82
83    pub fn get_committed_trace(
84        trace: DeviceMatrix<F>,
85        device: &GpuDevice,
86    ) -> CommittedTraceData<GpuBackend> {
87        let (commitment, data) = TraceCommitter::<GpuBackend>::commit(device, &[&trace]).unwrap();
88        CommittedTraceData {
89            commitment,
90            data: Arc::new(data),
91            trace,
92        }
93    }
94}
95
96impl Default for ProgramChipGPU {
97    fn default() -> Self {
98        panic!("ProgramChipGPU requires an explicit GpuDeviceCtx")
99    }
100}
101
102impl Chip<Vec<u32>, GpuBackend> for ProgramChipGPU {
103    fn generate_proving_ctx(&self, filtered_exec_freqs: Vec<u32>) -> AirProvingContext<GpuBackend> {
104        let cached = self.cached.clone().expect("Cached program must be loaded");
105        let height = cached.height();
106        let filtered_len = filtered_exec_freqs.len();
107        assert!(
108            filtered_len <= height,
109            "filtered_exec_freqs len={filtered_len} > cached trace height={height}"
110        );
111        let mut buffer: DeviceBuffer<F> = DeviceBuffer::with_capacity_on(height, &self.device_ctx);
112
113        filtered_exec_freqs
114            .into_iter()
115            .map(F::from_u32)
116            .collect::<Vec<_>>()
117            .copy_to_on(&mut buffer, &self.device_ctx)
118            .unwrap();
119        // Making sure to zero-out the untouched part of the buffer.
120        if filtered_len < height {
121            buffer
122                .fill_zero_suffix_on(filtered_len, &self.device_ctx)
123                .unwrap();
124        }
125
126        let common_main = DeviceMatrix::new(Arc::new(buffer), height, 1);
127
128        AirProvingContext {
129            cached_mains: vec![cached],
130            common_main,
131            public_values: vec![],
132        }
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use std::sync::Arc;
139
140    use openvm_cuda_backend::{data_transporter::assert_eq_host_and_device_matrix, prelude::F};
141    use openvm_instructions::{
142        instruction::Instruction,
143        program::{Program, DEFAULT_PC_STEP},
144        LocalOpcode,
145        SystemOpcode::*,
146    };
147    use openvm_stark_backend::{prover::TraceCommitter, StarkEngine};
148
149    use super::ProgramChipGPU;
150    use crate::{
151        system::program::{
152            tests::{BEQ, BNE, JAL, STOREW, SUB},
153            trace::generate_cached_trace,
154        },
155        utils::{test_cpu_engine, test_gpu_engine},
156    };
157
158    fn test_cached_committed_trace_data(program: Program<F>) {
159        let gpu_engine = test_gpu_engine();
160        let gpu_device = gpu_engine.device();
161        let gpu_trace =
162            ProgramChipGPU::generate_cached_trace(program.clone(), &gpu_device.device_ctx);
163        let gpu_cached = ProgramChipGPU::get_committed_trace(gpu_trace, gpu_device);
164
165        let cpu_engine = test_cpu_engine();
166        let cpu_device = cpu_engine.device();
167        let cpu_trace = Arc::new(generate_cached_trace(&program));
168        let (cpu_commit, _) = cpu_device.commit(&[&cpu_trace]).unwrap();
169
170        // NOTE: This compares the stacked matrices, not the original cached trace
171        assert_eq_host_and_device_matrix(cpu_trace, &gpu_cached.trace, &gpu_device.device_ctx);
172        assert_eq!(gpu_cached.commitment, cpu_commit);
173    }
174
175    #[test]
176    fn test_cuda_program_cached_tracegen_1() {
177        let instructions = vec![
178            Instruction::large_from_isize(STOREW, 2, 0, 0, 0, 1, 0, 1),
179            Instruction::large_from_isize(STOREW, 1, 1, 0, 0, 1, 0, 1),
180            Instruction::from_isize(BEQ, 0, 0, 3 * DEFAULT_PC_STEP as isize, 1, 0),
181            Instruction::from_isize(SUB, 0, 0, 1, 1, 1),
182            Instruction::from_isize(JAL, 2, -2 * (DEFAULT_PC_STEP as isize), 0, 1, 0),
183            Instruction::from_isize(TERMINATE.global_opcode(), 0, 0, 0, 0, 0),
184        ];
185        let program = Program::from_instructions(&instructions);
186        test_cached_committed_trace_data(program);
187    }
188
189    #[test]
190    fn test_cuda_program_cached_tracegen_2() {
191        let instructions = vec![
192            Instruction::large_from_isize(STOREW, 5, 0, 0, 0, 1, 0, 1),
193            Instruction::from_isize(BNE, 0, 4, 3 * DEFAULT_PC_STEP as isize, 1, 0),
194            Instruction::from_isize(JAL, 2, -2 * DEFAULT_PC_STEP as isize, 0, 1, 0),
195            Instruction::from_isize(TERMINATE.global_opcode(), 0, 0, 0, 0, 0),
196            Instruction::from_isize(BEQ, 0, 5, -(DEFAULT_PC_STEP as isize), 1, 0),
197        ];
198        let program = Program::from_instructions(&instructions);
199        test_cached_committed_trace_data(program);
200    }
201
202    #[test]
203    fn test_cuda_program_cached_tracegen_undefined_instructions() {
204        let instructions = vec![
205            Some(Instruction::large_from_isize(STOREW, 2, 0, 0, 0, 1, 0, 1)),
206            Some(Instruction::large_from_isize(STOREW, 1, 1, 0, 0, 1, 0, 1)),
207            Some(Instruction::from_isize(
208                BEQ,
209                0,
210                2,
211                3 * DEFAULT_PC_STEP as isize,
212                1,
213                0,
214            )),
215            None,
216            None,
217            Some(Instruction::from_isize(
218                TERMINATE.global_opcode(),
219                0,
220                0,
221                0,
222                0,
223                0,
224            )),
225        ];
226        let program = Program::new_without_debug_infos_with_option(&instructions, 0);
227        test_cached_committed_trace_data(program);
228    }
229}
openvm_circuit/system/cuda/program.rs

openvm_circuit/system/cuda/
program.rs
