changes so far

Author: qwang98

Cargo.lock

@@ -275,9 +275,9 @@ cuda-runtime-sys = "0.3.0-alpha.1"
 ignored = ["cargo-openvm"]
 
 # # The local openvm also needs to have stark-backend patched so all types match.
-# [patch."https://github.com/powdr-labs/stark-backend.git"]
-# openvm-stark-backend = { path = "../stark-backend/crates/stark-backend", default-features = false }
-# openvm-stark-sdk = { path = "../stark-backend/crates/stark-sdk", default-features = false }
-# openvm-cuda-backend = { path = "../stark-backend/crates/cuda-backend", default-features = false }
-# openvm-cuda-builder = { path = "../stark-backend/crates/cuda-builder", default-features = false }
-# openvm-cuda-common = { path = "../stark-backend/crates/cuda-common", default-features = false }
+[patch."https://github.com/powdr-labs/stark-backend.git"]
+openvm-stark-backend = { path = "../stark-backend/crates/stark-backend", default-features = false }
+openvm-stark-sdk = { path = "../stark-backend/crates/stark-sdk", default-features = false }
+openvm-cuda-backend = { path = "../stark-backend/crates/cuda-backend", default-features = false }
+openvm-cuda-builder = { path = "../stark-backend/crates/cuda-builder", default-features = false }
+openvm-cuda-common = { path = "../stark-backend/crates/cuda-common", default-features = false }

Cargo.toml

@@ -34,7 +34,7 @@ __global__ void alu_tracegen(
     // Fp *d_apc_trace,
     uint32_t *subs, // same length as dummy width
     uint32_t calls_per_apc_row, // 1 for non-apc
-    size_t width // dummy width
+    size_t width // dummy width or apc width
     // uint32_t *apc_row_index, // dummy row mapping to apc row same length as d_records
 ) {
     uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
@@ -56,7 +56,12 @@ __global__ void alu_tracegen(
         core.fill_trace_row_new(row.slice_from(COL_INDEX(Rv32BaseAluCols, core), number_of_gaps_in(sub, sizeof(Rv32BaseAluCols<uint8_t>))), rec.core, sub);
     } else {
         // TODO: use APC width if APC
-        row.fill_zero(0, sizeof(Rv32BaseAluCols<uint8_t>));
+        // this is now a hack because calls_per_apc_row can still be 1 even if we are in an APC
+        if (calls_per_apc_row == 1) {
+            row.fill_zero(0, sizeof(Rv32BaseAluCols<uint8_t>));
+        } else {
+            row.fill_zero(0, width);
+        }
     }
 }
 
@@ -77,7 +82,7 @@ extern "C" int _alu_tracegen(
 ) {
     assert((height & (height - 1)) == 0);
     assert(height >= d_records.len());
-    assert(width == sizeof(Rv32BaseAluCols<uint8_t>));
+    // assert(width == sizeof(Rv32BaseAluCols<uint8_t>)); // this is no longer true for APC
     auto [grid, block] = kernel_launch_params(height);
     alu_tracegen<<<grid, block>>>(
         d_trace,
@@ -91,7 +96,7 @@ extern "C" int _alu_tracegen(
         // Fp *d_apc_trace,
         subs, // same length as dummy width
         calls_per_apc_row, // 1 for non-apc
-        width // dummy width
+        width // dummy width or apc width
         // uint32_t *apc_row_index, // dummy row mapping to apc row same length as d_records
     );
     return CHECK_KERNEL();

extensions/rv32im/circuit/cuda/src/alu.cu

@@ -9,6 +9,7 @@ use openvm_cuda_backend::{
     base::DeviceMatrix, chip::get_empty_air_proving_ctx, prover_backend::GpuBackend, types::F,
 };
 use openvm_cuda_common::copy::MemCopyH2D;
+use openvm_cuda_common::d_buffer::DeviceBuffer;
 use openvm_stark_backend::{prover::types::AirProvingContext, Chip};
 
 use crate::{
@@ -27,6 +28,44 @@ pub struct Rv32BaseAluChipGpu {
 }
 
 impl Chip<DenseRecordArena, GpuBackend> for Rv32BaseAluChipGpu {
+    fn  generate_proving_ctx_new(&self, arena: DenseRecordArena, d_trace: &DeviceBuffer<F>, d_subs: &DeviceBuffer<u32>, calls_per_apc_row: u32) {
+        const RECORD_SIZE: usize = size_of::<(
+            Rv32BaseAluAdapterRecord,
+            BaseAluCoreRecord<RV32_REGISTER_NUM_LIMBS>,
+        )>();
+        let records = arena.allocated();
+        if records.is_empty() {
+            return;
+            // return get_empty_air_proving_ctx::<GpuBackend>();
+        }
+        debug_assert_eq!(records.len() % RECORD_SIZE, 0);
+
+        let trace_width = BaseAluCoreCols::<F, RV32_REGISTER_NUM_LIMBS, RV32_CELL_BITS>::width()
+            + Rv32BaseAluAdapterCols::<F>::width();
+        let trace_height = next_power_of_two_or_zero(records.len() / RECORD_SIZE);
+
+        let d_records = records.to_device().unwrap();
+
+        unsafe {
+            tracegen(
+                d_trace, // replaced with apc trace
+                trace_height,
+                &d_records,
+                &self.range_checker.count,
+                self.range_checker.count.len(),
+                &self.bitwise_lookup.count,
+                RV32_CELL_BITS,
+                self.timestamp_max_bits as u32,
+                // d_apc_trace.buffer(),
+                d_subs, // same length as dummy width
+                calls_per_apc_row,
+            )
+            .unwrap();
+        }
+
+        // AirProvingContext::simple_no_pis(d_trace)
+    }
+
     fn generate_proving_ctx(&self, arena: DenseRecordArena) -> AirProvingContext<GpuBackend> {
         const RECORD_SIZE: usize = size_of::<(
             Rv32BaseAluAdapterRecord,