feat: RISC-V 256-bit integer branch chip adapter + core implementatio…

…ns (#747)

* feat: RISC-V 256-bit integer branch chip adapter + core implementations

* Feat: add range checks to vec adapter (#746)

* feat: add range checks to vec_heap_adapter

* address comments

* feat: add range checks to memory pointer

* merge main

* fix: assert

* chore: move range check from execute to generate_trace

---------

Co-authored-by: Arayi Khalatyan <[email protected]>
Co-authored-by: Arayi <[email protected]>
Co-authored-by: Jonathan Wang <[email protected]>

toolchain/instructions/src/lib.rs

@@ -330,6 +330,14 @@ pub struct Rv32Shift256Opcode(pub ShiftOpcode);
 #[opcode_offset = 0x408]
 pub struct Rv32LessThan256Opcode(pub LessThanOpcode);
 
+#[derive(Copy, Clone, Debug, UsizeOpcode)]
+#[opcode_offset = 0x420]
+pub struct Rv32BranchEqual256Opcode(pub BranchEqualOpcode);
+
+#[derive(Copy, Clone, Debug, UsizeOpcode)]
+#[opcode_offset = 0x425]
+pub struct Rv32BranchLessThan256Opcode(pub BranchLessThanOpcode);
+
 #[derive(Copy, Clone, Debug, UsizeOpcode)]
 #[opcode_offset = 0x450]
 pub struct Rv32Mul256Opcode(pub MulOpcode);

vm/src/arch/chip_set.rs

@@ -7,7 +7,7 @@ use std::{
     sync::Arc,
 };
 
-use adapters::Rv32HeapAdapterChip;
+use adapters::{Rv32HeapAdapterChip, Rv32HeapBranchAdapterChip};
 use ax_circuit_primitives::{
     bitwise_op_lookup::{BitwiseOperationLookupBus, BitwiseOperationLookupChip},
     range_tuple::{RangeTupleCheckerBus, RangeTupleCheckerChip},
@@ -43,7 +43,8 @@ use crate::{
         },
         hashes::{keccak256::KeccakVmChip, poseidon2::Poseidon2Chip},
         int256::{
-            Rv32BaseAlu256Chip, Rv32LessThan256Chip, Rv32Multiplication256Chip, Rv32Shift256Chip,
+            Rv32BaseAlu256Chip, Rv32BranchEqual256Chip, Rv32BranchLessThan256Chip,
+            Rv32LessThan256Chip, Rv32Multiplication256Chip, Rv32Shift256Chip,
         },
         modular::{
             ModularAddSubChip, ModularAddSubCoreChip, ModularMulDivChip, ModularMulDivCoreChip,
@@ -710,6 +711,38 @@ impl VmConfig {
                     }
                     chips.push(AxVmChip::Shift256Rv32(chip));
                 }
+                ExecutorName::BranchEqual256Rv32 => {
+                    let chip = Rc::new(RefCell::new(Rv32BranchEqual256Chip::new(
+                        Rv32HeapBranchAdapterChip::new(
+                            execution_bus,
+                            program_bus,
+                            memory_controller.clone(),
+                            bitwise_lookup_chip.clone(),
+                        ),
+                        BranchEqualCoreChip::new(offset, DEFAULT_PC_STEP),
+                        memory_controller.clone(),
+                    )));
+                    for opcode in range {
+                        executors.insert(opcode, chip.clone().into());
+                    }
+                    chips.push(AxVmChip::BranchEqual256Rv32(chip));
+                }
+                ExecutorName::BranchLessThan256Rv32 => {
+                    let chip = Rc::new(RefCell::new(Rv32BranchLessThan256Chip::new(
+                        Rv32HeapBranchAdapterChip::new(
+                            execution_bus,
+                            program_bus,
+                            memory_controller.clone(),
+                            bitwise_lookup_chip.clone(),
+                        ),
+                        BranchLessThanCoreChip::new(bitwise_lookup_chip.clone(), offset),
+                        memory_controller.clone(),
+                    )));
+                    for opcode in range {
+                        executors.insert(opcode, chip.clone().into());
+                    }
+                    chips.push(AxVmChip::BranchLessThan256Rv32(chip));
+                }
                 ExecutorName::CastF => {
                     let chip = Rc::new(RefCell::new(CastFChip::new(
                         ConvertAdapterChip::new(
@@ -1317,21 +1350,11 @@ fn default_executor_range(executor: ExecutorName) -> (Range<usize>, usize) {
             Rv32AuipcOpcode::COUNT,
             Rv32AuipcOpcode::default_offset(),
         ),
-        ExecutorName::BaseAlu256Rv32 => (
-            Rv32BaseAlu256Opcode::default_offset(),
-            BaseAluOpcode::COUNT,
-            Rv32BaseAlu256Opcode::default_offset(),
-        ),
         ExecutorName::LessThanRv32 => (
             LessThanOpcode::default_offset(),
             LessThanOpcode::COUNT,
             LessThanOpcode::default_offset(),
         ),
-        ExecutorName::LessThan256Rv32 => (
-            Rv32LessThan256Opcode::default_offset(),
-            LessThanOpcode::COUNT,
-            Rv32LessThan256Opcode::default_offset(),
-        ),
         ExecutorName::MultiplicationRv32 => (
             MulOpcode::default_offset(),
             MulOpcode::COUNT,
@@ -1342,11 +1365,6 @@ fn default_executor_range(executor: ExecutorName) -> (Range<usize>, usize) {
             MulHOpcode::COUNT,
             MulHOpcode::default_offset(),
         ),
-        ExecutorName::Multiplication256Rv32 => (
-            Rv32Mul256Opcode::default_offset(),
-            MulOpcode::COUNT,
-            Rv32Mul256Opcode::default_offset(),
-        ),
         ExecutorName::DivRemRv32 => (
             DivRemOpcode::default_offset(),
             DivRemOpcode::COUNT,
@@ -1357,11 +1375,6 @@ fn default_executor_range(executor: ExecutorName) -> (Range<usize>, usize) {
             ShiftOpcode::COUNT,
             ShiftOpcode::default_offset(),
         ),
-        ExecutorName::Shift256Rv32 => (
-            Rv32Shift256Opcode::default_offset(),
-            ShiftOpcode::COUNT,
-            Rv32Shift256Opcode::default_offset(),
-        ),
         ExecutorName::BranchEqualRv32 => (
             BranchEqualOpcode::default_offset(),
             BranchEqualOpcode::COUNT,
@@ -1372,6 +1385,36 @@ fn default_executor_range(executor: ExecutorName) -> (Range<usize>, usize) {
             BranchLessThanOpcode::COUNT,
             BranchLessThanOpcode::default_offset(),
         ),
+        ExecutorName::BaseAlu256Rv32 => (
+            Rv32BaseAlu256Opcode::default_offset(),
+            BaseAluOpcode::COUNT,
+            Rv32BaseAlu256Opcode::default_offset(),
+        ),
+        ExecutorName::LessThan256Rv32 => (
+            Rv32LessThan256Opcode::default_offset(),
+            LessThanOpcode::COUNT,
+            Rv32LessThan256Opcode::default_offset(),
+        ),
+        ExecutorName::Multiplication256Rv32 => (
+            Rv32Mul256Opcode::default_offset(),
+            MulOpcode::COUNT,
+            Rv32Mul256Opcode::default_offset(),
+        ),
+        ExecutorName::Shift256Rv32 => (
+            Rv32Shift256Opcode::default_offset(),
+            ShiftOpcode::COUNT,
+            Rv32Shift256Opcode::default_offset(),
+        ),
+        ExecutorName::BranchEqual256Rv32 => (
+            Rv32BranchEqual256Opcode::default_offset(),
+            BranchEqualOpcode::COUNT,
+            Rv32BranchEqual256Opcode::default_offset(),
+        ),
+        ExecutorName::BranchLessThan256Rv32 => (
+            Rv32BranchLessThan256Opcode::default_offset(),
+            BranchLessThanOpcode::COUNT,
+            Rv32BranchLessThan256Opcode::default_offset(),
+        ),
         ExecutorName::CastF => (
             CastfOpcode::default_offset(),
             CastfOpcode::COUNT,

vm/src/arch/chips.rs

@@ -24,7 +24,8 @@ use crate::{
         },
         hashes::{keccak256::KeccakVmChip, poseidon2::Poseidon2Chip},
         int256::{
-            Rv32BaseAlu256Chip, Rv32LessThan256Chip, Rv32Multiplication256Chip, Rv32Shift256Chip,
+            Rv32BaseAlu256Chip, Rv32BranchEqual256Chip, Rv32BranchLessThan256Chip,
+            Rv32LessThan256Chip, Rv32Multiplication256Chip, Rv32Shift256Chip,
         },
         modular::{ModularAddSubChip, ModularMulDivChip},
     },
@@ -105,6 +106,8 @@ pub enum AxVmInstructionExecutor<F: PrimeField32> {
     BaseAlu256Rv32(Rc<RefCell<Rv32BaseAlu256Chip<F>>>),
     Shift256Rv32(Rc<RefCell<Rv32Shift256Chip<F>>>),
     LessThan256Rv32(Rc<RefCell<Rv32LessThan256Chip<F>>>),
+    BranchEqual256Rv32(Rc<RefCell<Rv32BranchEqual256Chip<F>>>),
+    BranchLessThan256Rv32(Rc<RefCell<Rv32BranchLessThan256Chip<F>>>),
     Multiplication256Rv32(Rc<RefCell<Rv32Multiplication256Chip<F>>>),
     // Modular arithmetic:
     ModularAddSubRv32_1x32(Rc<RefCell<ModularAddSubChip<F, 1, 32>>>),
@@ -165,6 +168,8 @@ pub enum AxVmChip<F: PrimeField32> {
     LessThan256Rv32(Rc<RefCell<Rv32LessThan256Chip<F>>>),
     Multiplication256Rv32(Rc<RefCell<Rv32Multiplication256Chip<F>>>),
     Shift256Rv32(Rc<RefCell<Rv32Shift256Chip<F>>>),
+    BranchEqual256Rv32(Rc<RefCell<Rv32BranchEqual256Chip<F>>>),
+    BranchLessThan256Rv32(Rc<RefCell<Rv32BranchLessThan256Chip<F>>>),
     // Modular arithmetic:
     ModularAddSubRv32_1x32(Rc<RefCell<ModularAddSubChip<F, 1, 32>>>),
     ModularMulDivRv32_1x32(Rc<RefCell<ModularMulDivChip<F, 1, 32>>>),

vm/src/arch/config.rs

@@ -115,6 +115,11 @@ impl VmConfig {
             .add_executor(ExecutorName::Shift256Rv32)
     }
 
+    pub fn add_int256_branch(self) -> Self {
+        self.add_executor(ExecutorName::BranchEqual256Rv32)
+            .add_executor(ExecutorName::BranchLessThan256Rv32)
+    }
+
     pub fn add_int256_m(self) -> Self {
         self.add_executor(ExecutorName::Multiplication256Rv32)
     }

vm/src/intrinsics/int256/mod.rs

@@ -1,8 +1,11 @@
 use crate::{
     arch::VmChipWrapper,
     rv32im::{
-        adapters::{Rv32HeapAdapterChip, INT256_NUM_LIMBS, RV32_CELL_BITS},
-        BaseAluCoreChip, LessThanCoreChip, MultiplicationCoreChip, ShiftCoreChip,
+        adapters::{
+            Rv32HeapAdapterChip, Rv32HeapBranchAdapterChip, INT256_NUM_LIMBS, RV32_CELL_BITS,
+        },
+        BaseAluCoreChip, BranchEqualCoreChip, BranchLessThanCoreChip, LessThanCoreChip,
+        MultiplicationCoreChip, ShiftCoreChip,
     },
 };
 
@@ -32,3 +35,15 @@ pub type Rv32Shift256Chip<F> = VmChipWrapper<
     Rv32HeapAdapterChip<F, 2, INT256_NUM_LIMBS, INT256_NUM_LIMBS>,
     ShiftCoreChip<INT256_NUM_LIMBS, RV32_CELL_BITS>,
 >;
+
+pub type Rv32BranchEqual256Chip<F> = VmChipWrapper<
+    F,
+    Rv32HeapBranchAdapterChip<F, 2, INT256_NUM_LIMBS>,
+    BranchEqualCoreChip<INT256_NUM_LIMBS>,
+>;
+
+pub type Rv32BranchLessThan256Chip<F> = VmChipWrapper<
+    F,
+    Rv32HeapBranchAdapterChip<F, 2, INT256_NUM_LIMBS>,
+    BranchLessThanCoreChip<INT256_NUM_LIMBS, RV32_CELL_BITS>,
+>;