added offset support for APC and sub; currently backward compatible with non-APC but not sure if APC works

Author: qwang98

crates/circuits/primitives/cuda/include/primitives/histogram.cuh

@@ -86,7 +86,7 @@ struct VariableRangeChecker {
     __device__ __forceinline__ void decompose_new(
         uint32_t x,
         size_t bits,
-        RowSlice limbs,
+        RowSliceNew limbs,
         const size_t limbs_len,
         uint32_t *sub
     ) {

crates/circuits/primitives/cuda/include/primitives/less_than.cuh

@@ -47,7 +47,7 @@ __device__ __forceinline__ void generate_subrow_new(
     uint32_t x,
     uint32_t y,
     const size_t lower_decomp_len,
-    RowSlice lower_decomp,
+    RowSliceNew lower_decomp,
     uint32_t *sub
 ) {
     rc.decompose_new(y - x - 1, max_bits, lower_decomp, lower_decomp_len, sub);

crates/circuits/primitives/cuda/include/primitives/trace_access.h

@@ -4,6 +4,79 @@
 #include <cstddef>
 #include <cstdint>
 
+
+/// A RowSlice is a contiguous section of a row in col-based trace.
+struct RowSliceNew {
+    Fp *ptr;
+    size_t stride;
+    size_t optimized_offset;
+    size_t dummy_offset;
+
+
+    __device__ RowSliceNew(Fp *ptr, size_t stride, size_t optimized_offset, size_t dummy_offset) : ptr(ptr), stride(stride), optimized_offset(optimized_offset), dummy_offset(dummy_offset) {}
+
+    __device__ __forceinline__ Fp &operator[](size_t column_index) const {
+        // While implementing tracegen for SHA256, we encountered what we believe to be an nvcc
+        // compiler bug. Occasionally, at various non-zero PTXAS optimization levels the compiler
+        // tries to replace this multiplication with a series of SHL, ADD, and AND instructions
+        // that we believe erroneously adds ~2^49 to the final address via an improper carry
+        // propagation. To read more, see https://github.com/stephenh-axiom-xyz/cuda-illegal.
+        return ptr[column_index * stride];
+    }
+
+    __device__ static RowSliceNew null() { return RowSliceNew(nullptr, 0, 0, 0); }
+
+    __device__ bool is_valid() const { return ptr != nullptr; }
+
+    template <typename T>
+    __device__ __forceinline__ void write(size_t column_index, T value) const {
+        ptr[column_index * stride] = value;
+    }
+
+    template <typename T>
+    __device__ __forceinline__ void write_array(size_t column_index, size_t length, const T *values)
+        const {
+#pragma unroll
+        for (size_t i = 0; i < length; i++) {
+            ptr[(column_index + i) * stride] = values[i];
+        }
+    }
+
+    template <typename T>
+    __device__ __forceinline__ void write_array_new(size_t column_index, size_t length, const T *values, const uint32_t *sub)
+        const {
+#pragma unroll
+        for (size_t i = 0; i < length; i++) {
+            if (sub[i] != UINT32_MAX) {
+                ptr[(column_index + i) * stride] = values[i];
+            }
+        }
+    }
+
+    template <typename T>
+    __device__ __forceinline__ void write_bits(size_t column_index, const T value) const {
+#pragma unroll
+        for (size_t i = 0; i < sizeof(T) * 8; i++) {
+            ptr[(column_index + i) * stride] = (value >> i) & 1;
+        }
+    }
+
+    __device__ __forceinline__ void fill_zero(size_t column_index_from, size_t length) const {
+#pragma unroll
+        for (size_t i = 0, c = column_index_from; i < length; i++, c++) {
+            ptr[c * stride] = 0;
+        }
+    }
+
+    __device__ __forceinline__ RowSliceNew slice_from(size_t column_index, uint32_t gap) const {
+        return RowSliceNew(ptr + (column_index - gap) * stride, stride, optimized_offset + column_index - gap, dummy_offset + column_index);
+    }
+
+    __device__ __forceinline__ RowSliceNew shift_row(size_t n) const {
+        return RowSliceNew(ptr + n, stride, optimized_offset, dummy_offset);
+    }
+};
+
 /// A RowSlice is a contiguous section of a row in col-based trace.
 struct RowSlice {
     Fp *ptr;
@@ -87,9 +160,9 @@ struct RowSlice {
 #define COL_WRITE_VALUE_NEW(ROW, STRUCT, FIELD, VALUE, SUB)                                         \
     do {                                                                                            \
         const size_t _col_idx = COL_INDEX(STRUCT, FIELD);                                           \
-        const auto _apc_idx = (SUB)[_col_idx];                                                      \
+        const auto _apc_idx = (SUB)[_col_idx + ROW.dummy_offset];                                                      \
         if (_apc_idx != UINT32_MAX) {                                                               \
-            (ROW).write(_apc_idx, VALUE);                                                           \
+            (ROW).write(_apc_idx - ROW.optimized_offset, VALUE);                                                           \
         }                                                                                           \
     } while (0)
 

crates/vm/cuda/include/system/memory/controller.cuh

@@ -23,14 +23,14 @@ struct MemoryAuxColsFactory {
         COL_WRITE_VALUE(row, MemoryBaseAuxCols, prev_timestamp, prev_timestamp);
     }
 
-    __device__ void fill_new(RowSlice row, uint32_t prev_timestamp, uint32_t timestamp, uint32_t *sub) {
+    __device__ void fill_new(RowSliceNew row, uint32_t prev_timestamp, uint32_t timestamp, uint32_t *sub) {
         AssertLessThan::generate_subrow_new(
             range_checker,
             timestamp_max_bits,
             prev_timestamp,
             timestamp,
             AUX_LEN,
-            row.slice_from(COL_INDEX(MemoryBaseAuxCols, timestamp_lt_aux) - number_of_gaps_in(sub, sizeof(MemoryBaseAuxCols<uint8_t>))),
+            row.slice_from(COL_INDEX(MemoryBaseAuxCols, timestamp_lt_aux), number_of_gaps_in(sub, sizeof(MemoryBaseAuxCols<uint8_t>))),
             sub
         );
         COL_WRITE_VALUE_NEW(row, MemoryBaseAuxCols, prev_timestamp, prev_timestamp, sub);

extensions/rv32im/circuit/cuda/include/rv32im/adapters/alu.cuh

@@ -82,7 +82,7 @@ struct Rv32BaseAluAdapter {
         );
     }
 
-    __device__ void fill_trace_row_new(RowSlice row, Rv32BaseAluAdapterRecord record, uint32_t *sub) {
+    __device__ void fill_trace_row_new(RowSliceNew row, Rv32BaseAluAdapterRecord record, uint32_t *sub) {
         COL_WRITE_VALUE_NEW(row, Rv32BaseAluAdapterCols, from_state.pc, record.from_pc, sub);
         COL_WRITE_VALUE_NEW(row, Rv32BaseAluAdapterCols, from_state.timestamp, record.from_timestamp, sub);
 
@@ -93,7 +93,7 @@ struct Rv32BaseAluAdapter {
 
         // Read auxiliary for rs1
         mem_helper.fill_new(
-            row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[0]) - number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
+            row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[0]), number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
             record.reads_aux[0].prev_timestamp,
             record.from_timestamp,
             sub
@@ -102,13 +102,13 @@ struct Rv32BaseAluAdapter {
         // rs2: register read when rs2_as == RV32_REGISTER_AS (== 1), otherwise immediate.
         if (record.rs2_as != 0) {
             mem_helper.fill_new(
-                row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[1]) - number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
+                row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[1]), number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
                 record.reads_aux[1].prev_timestamp,
                 record.from_timestamp + 1,
                 sub
             );
         } else {
-            RowSlice rs2_aux = row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[1]) - number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>)));
+            RowSliceNew rs2_aux = row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, reads_aux[1]), number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>)));
 #pragma unroll
             for (size_t i = 0; i < sizeof(MemoryReadAuxCols<uint8_t>); i++) {
                 rs2_aux.write(i, 0);
@@ -117,11 +117,11 @@ struct Rv32BaseAluAdapter {
             bitwise_lookup.add_range(record.rs2 & mask, (record.rs2 >> RV32_CELL_BITS) & mask);
         }
 
-        COL_WRITE_ARRAY(
-            row, Rv32BaseAluAdapterCols, writes_aux.prev_data, record.writes_aux.prev_data
+        COL_WRITE_ARRAY_NEW(
+            row, Rv32BaseAluAdapterCols, writes_aux.prev_data, record.writes_aux.prev_data, sub
         );
         mem_helper.fill_new(
-            row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, writes_aux) - number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
+            row.slice_from(COL_INDEX(Rv32BaseAluAdapterCols, writes_aux), number_of_gaps_in(sub, sizeof(Rv32BaseAluAdapterCols<uint8_t>))),
             record.writes_aux.prev_timestamp,
             record.from_timestamp + 2,
             sub

extensions/rv32im/circuit/cuda/include/rv32im/cores/alu.cuh

@@ -140,7 +140,7 @@ template <size_t NUM_LIMBS> struct BaseAluCore {
         }
     }
 
-    __device__ void fill_trace_row_new(RowSlice row, BaseAluCoreRecord<NUM_LIMBS> record, uint32_t *sub) {
+    __device__ void fill_trace_row_new(RowSliceNew row, BaseAluCoreRecord<NUM_LIMBS> record, uint32_t *sub) {
         uint8_t a[NUM_LIMBS];
         uint8_t carry_buf[NUM_LIMBS];
 

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

@@ -23,26 +23,27 @@ struct Rv32BaseAluRecord {
 };
 
 __global__ void alu_tracegen(
-    Fp *d_trace,
-    size_t height,
+    Fp *d_trace, // can be apc trace
+    size_t height, // can be apc height
     DeviceBufferConstView<Rv32BaseAluRecord> d_records,
     uint32_t *d_range_checker_ptr,
     size_t range_checker_bins,
     uint32_t *d_bitwise_lookup_ptr,
     size_t bitwise_num_bits,
     uint32_t timestamp_max_bits,
     // Fp *d_apc_trace,
-    uint32_t *subs // same length as dummy width
-    // size_t width, // dummy width
+    uint32_t *subs, // same length as dummy width
+    uint32_t calls_per_apc_row, // 1 for non-apc
+    size_t width // dummy 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;
-    RowSlice row(d_trace + idx, height);
+    RowSliceNew row(d_trace + idx / calls_per_apc_row, height, 0, 0); // we need to slice to the correct APC row, but if non-APC it's dividing by 1 and therefore the same idx
     if (idx < d_records.len()) {
         auto const &rec = d_records[idx];
         // RowSlice apc_row(d_apc_trace + apc_row_index[idx], height);
         // auto const sub = subs[idx * width]; // offset the subs to the corresponding dummy row
-        uint32_t *sub = subs;
+        uint32_t *sub = &subs[(idx % calls_per_apc_row) * width]; // dummy width
 
         Rv32BaseAluAdapter adapter(
             VariableRangeChecker(d_range_checker_ptr, range_checker_bins),
@@ -52,7 +53,7 @@ __global__ void alu_tracegen(
         adapter.fill_trace_row_new(row, rec.adapter, sub);
 
         Rv32BaseAluCore core(BitwiseOperationLookup(d_bitwise_lookup_ptr, bitwise_num_bits));
-        core.fill_trace_row_new(row.slice_from(COL_INDEX(Rv32BaseAluCols, core) - number_of_gaps_in(sub, sizeof(Rv32BaseAluCols<uint8_t>))), rec.core, sub);
+        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>));
@@ -70,7 +71,8 @@ extern "C" int _alu_tracegen(
     size_t bitwise_num_bits,
     uint32_t timestamp_max_bits,
     // Fp *d_apc_trace,
-    uint32_t *subs // same length as dummy width
+    uint32_t *subs, // same length as dummy width
+    uint32_t calls_per_apc_row // 1 for non-apc
     // uint32_t *apc_row_index, // dummy row mapping to apc row same length as d_records
 ) {
     assert((height & (height - 1)) == 0);
@@ -87,8 +89,9 @@ extern "C" int _alu_tracegen(
         bitwise_num_bits,
         timestamp_max_bits,
         // Fp *d_apc_trace,
-        subs // same length as dummy width
-        // size_t width, // dummy width
+        subs, // same length as dummy width
+        calls_per_apc_row, // 1 for non-apc
+        width // dummy width
         // uint32_t *apc_row_index, // dummy row mapping to apc row same length as d_records
     );
     return CHECK_KERNEL();

extensions/rv32im/circuit/src/base_alu/cuda.rs

@@ -63,7 +63,7 @@ impl Chip<DenseRecordArena, GpuBackend> for Rv32BaseAluChipGpu {
                 self.timestamp_max_bits as u32,
                 // d_apc_trace.buffer(),
                 &d_subs, // same length as dummy width
-                         // apc_row_index, // dummy row mapping to apc row same length as d_records
+                1, // calls_per_apc_row: 1 for non-apc
             )
             .unwrap();
         }

extensions/rv32im/circuit/src/cuda_abi.rs

@@ -324,7 +324,7 @@ pub mod alu_cuda {
             timestamp_max_bits: u32,
             // d_apc_trace: *mut F,
             d_subs: *mut u32,
-            // apc_row_index: *mut u32,
+            calls_per_apc_row: u32, // 1 for non-apc
         ) -> i32;
     }
 
@@ -339,7 +339,7 @@ pub mod alu_cuda {
         timestamp_max_bits: u32,
         // d_apc_trace: &DeviceBuffer<F>,
         d_subs: &DeviceBuffer<u32>,
-        // apc_row_index: Option<Vec<u32>>,
+        calls_per_apc_row: u32 // 1 for non-apc
     ) -> Result<(), CudaError> {
         let width = d_trace.len() / height;
         CudaError::from_result(_alu_tracegen(
@@ -354,6 +354,7 @@ pub mod alu_cuda {
             timestamp_max_bits,
             // d_apc_trace.as_mut_ptr(),
             d_subs.as_mut_ptr() as *mut u32,
+            calls_per_apc_row, // 1 for non-apc
             // apc_row_index,
         ))
     }