Finish M extension

.github/workflows/rust.yml

@@ -49,7 +49,7 @@ jobs:
       - name: Cache Jolt RISC-V Rust toolchain
         uses: actions/cache@v4
         with:
-          key: jolt-rust-toolchain-${{hashFiles('.jolt.rust.toolchain-tag')}}
+          key: jolt-rust-toolchain-${{hashFiles('guest-toolchain-tag')}}
           path: ~/.jolt
       - name: Install Jolt RISC-V Rust toolchain
         run: cargo run install-toolchain

Cargo.toml

@@ -44,6 +44,8 @@ members = [
     "examples/alloc/guest",
     "examples/stdlib",
     "examples/stdlib/guest",
+    "examples/muldiv",
+    "examples/muldiv/guest",
 ]
 
 [features]

book/src/SUMMARY.md

@@ -13,11 +13,12 @@
     - [Troubleshooting](./usage/troubleshooting.md)
 - [Contributors](./contributors.md)
     - [How it works](./how_it_works.md)
-        - [Jolt](./how/jolt.md)
+        - [Architecture overview](./how/architecture.md)
             - [Instruction lookups](./how/instruction_lookups.md)
             - [Read-write memory](./how/read_write_memory.md)
             - [Bytecode](./how/bytecode.md)
             - [R1CS constraints](./how/r1cs_constraints.md)
+        - [M extension](./how/m-extension.md)
     - [Background](./background.md)
         - [Sumcheck](./background/sumcheck.md)
         - [Multilinear Extensions](./background/multilinear-extensions.md)

book/src/dev/install.md

@@ -5,7 +5,7 @@
 - `curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh`
 - Rustup should automatically install Rust toolchain and necessary targets on
 the first `cargo` invocation.  If you need to add the RISC-V target for building
-guest programs manually use `rustup target add riscv32i-unknown-none-elf`.
+guest programs manually use `rustup target add riscv32im-unknown-none-elf`.
 
 ## mdBook
 

book/src/how/jolt.md → book/src/how/architecture.md

@@ -1,10 +1,10 @@
-# Jolt
+# Architecture overview
 
 This section gives an overview of the core components of Jolt.
 
 ## Jolt's four components
 
-A VM does two things: 
+A VM does two things:
 
 - Repeatedly execute the fetch-decode-execute logic of its instruction set architecture.
 - Perform reads and writes to Random Access Memory (RAM).
@@ -25,13 +25,13 @@ To handle reads/writes to RAM (and registers) Jolt uses a memory checking argume
 
 ### R1CS constraints
 
-To handle the "fetch" part of the fetch-decode-execute loop, there is a minimal R1CS instance (about 60 constraints per cycle of the RISC-V VM). These constraints handle program counter (PC) updates and serves as the "glue" enforcing consistency between polynomials used in the components below. Jolt uses [Spartan](https://eprint.iacr.org/2019/550), optimized for the highly-structured nature of the constraint system (e.g., the R1CS constraint matrices are block-diagonal with blocks of size only about 60 x 80). This is implemented in [jolt-core/src/r1cs](../../../jolt-core/src/r1cs/). 
+To handle the "fetch" part of the fetch-decode-execute loop, there is a minimal R1CS instance (about 60 constraints per cycle of the RISC-V VM). These constraints handle program counter (PC) updates and serves as the "glue" enforcing consistency between polynomials used in the components below. Jolt uses [Spartan](https://eprint.iacr.org/2019/550), optimized for the highly-structured nature of the constraint system (e.g., the R1CS constraint matrices are block-diagonal with blocks of size only about 60 x 80). This is implemented in [jolt-core/src/r1cs](../../../jolt-core/src/r1cs/).
 
 *For more details: [R1CS constraints](./r1cs_constraints.md)*
 
 ### Instruction lookups
 
-To handle the "execute" part of the fetch-decode-execute loop, Jolt invokes the Lasso lookup argument. The lookup argument maps every instruction (including its operands) to its output. This is implemented in [instruction_lookups.rs](https://github.com/a16z/jolt/blob/main/jolt-core/src/jolt/vm/instruction_lookups.rs). 
+To handle the "execute" part of the fetch-decode-execute loop, Jolt invokes the Lasso lookup argument. The lookup argument maps every instruction (including its operands) to its output. This is implemented in [instruction_lookups.rs](https://github.com/a16z/jolt/blob/main/jolt-core/src/jolt/vm/instruction_lookups.rs).
 
 *For more details: [Instruction lookups](./instruction_lookups.md)*
 

book/src/how/m-extension.md

@@ -0,0 +1,117 @@
+# M extension
+
+Jolt supports the RISC-V "M" extension for integer multiplication and division.
+The instructions included in this extension are described [here](https://msyksphinz-self.github.io/riscv-isadoc/html/rvm.html).
+For RV32, the M extension includes 8 instructions: `MUL`, `MULH`, `MULHSU`, `MULU`, `DIV`, `DIVU`, `REM`, and `REMU`.
+
+The [Jolt paper](https://eprint.iacr.org/2023/1217.pdf) describes how to handle the M extension instructions in Section 6,
+but our implementation deviates from the paper in a couple ways (described below).
+
+## Virtual sequences
+
+Section 6.1 of the Jolt paper introduces virtual instructions and registers –– some of the M extension
+instructions cannot be implemented as a single subtable decomposition, but rather must be split into
+a sequence of instructions which together compute the output and places it in the destination register.
+In our implementation, these sequences are captured by the `VirtualInstructionSequence` trait.
+
+The instructions that comprise such a sequence can be a combination of "real" RISC-V instructions and "virtual"
+instructions which only appear in the context of virtual sequences.
+We also introduce 32 virtual registers as "scratch space" where instructions in a virtual sequence
+can write intermediate values.
+
+## Deviations from the Jolt paper
+
+There are three inconsistencies between the virtual sequences provided in Section 6.3
+of the Jolt paper, and the RISC-V specification. Namely:
+
+1. The Jolt prover (as described in the paper) would fail to produce a valid proof
+if it encountered a division by zero; since the divisor `y` is 0, the `ASSERT_LTU`/`ASSERT_LT_ABS` would
+always fail (for `DIVU` and `DIV`, respectively).
+1. The MLE provided for `ASSERT_LT_ABS` in Section 6.1.1 doesn't account for two's complement.
+1. The `ASSERT_EQ_SIGNS` instruction should always return true if the remainder is 0.
+
+To address these issues, our implementation of `DIVU`, `DIV`, `REMU`, and `REM` deviate from the
+Jolt paper in the following ways.
+
+### `DIVU` virtual sequence
+
+1. `ADVICE` --, --, --, $v_q$   `// store non-deterministic advice` $q$ `into `$v_q$
+1. `ADVICE` --, --, --, $v_r$   `// store non-deterministic advice` $r$ `into `$v_r$
+1. `MUL` $v_q$, $r_y$, --, $v_{qy}$   `// compute q * y`
+1. `ASSERT_VALID_UNSIGNED_REMAINDER` $v_r$, $r_y$, --, --   `// assert that y == 0 || r < y`
+1. `ASSERT_LTE` $v_{qy}$, $r_x$, --, --   `// assert q * y <= x`
+1. `ASSERT_VALID_DIV0` $r_y$, $v_q$, --, --   `// assert that y != 0 || q == 2 ** WORD_SIZE - 1`
+1. `ADD` $v_{qy}$, $v_r$, --, $v_0$   `// compute q * y + r`
+1. `ASSERT_EQ` $v_0$, $x$, --, --
+1. `MOVE` $v_q$, --, --, `rd`
+
+### `REMU` virtual sequence
+
+1. `ADVICE` --, --, --, $v_q$   `// store non-deterministic advice` $q$ `into `$v_q$
+1. `ADVICE` --, --, --, $v_r$   `// store non-deterministic advice` $r$ `into `$v_r$
+1. `MUL` $v_q$, $r_y$, --, $v_{qy}$   `// compute q * y`
+1. `ASSERT_VALID_UNSIGNED_REMAINDER` $v_r$, $r_y$, --, --   `// assert that y == 0 || r < y`
+1. `ASSERT_LTE` $v_{qy}$, $r_x$, --, --   `// assert q * y <= x`
+1. `ADD` $v_{qy}$, $v_r$, --, $v_0$   `// compute q * y + r`
+1. `ASSERT_EQ` $v_0$, $x$, --, --
+1. `MOVE` $v_r$, --, --, `rd`
+
+### `DIV` virtual sequence
+
+1. `ADVICE` --, --, --, $v_q$   `// store non-deterministic advice` $q$ `into `$v_q$
+1. `ADVICE` --, --, --, $v_r$   `// store non-deterministic advice` $r$ `into `$v_r$
+1. `ASSERT_VALID_SIGNED_REMAINDER` $v_r$, $r_y$, --, --   `// assert that r == 0 || y == 0 || (|r| < |y| && sign(r) == sign(y))`
+1. `ASSERT_VALID_DIV0` $r_y$, $v_q$, --, --   `// assert that y != 0 || q == 2 ** WORD_SIZE - 1`
+1. `MUL` $v_q$, $r_y$, --, $v_{qy}$   `// compute q * y`
+1. `ADD` $v_{qy}$, $v_r$, --, $v_0$   `// compute q * y + r`
+1. `ASSERT_EQ` $v_0$, $x$, --, --
+1. `MOVE` $v_q$, --, --, `rd`
+
+### `REM` virtual sequence
+
+1. `ADVICE` --, --, --, $v_q$   `// store non-deterministic advice` $q$ `into `$v_q$
+1. `ADVICE` --, --, --, $v_r$   `// store non-deterministic advice` $r$ `into `$v_r$
+1. `ASSERT_VALID_SIGNED_REMAINDER` $v_r$, $r_y$, --, --   `// assert that r == 0 || y == 0 || (|r| < |y| && sign(r) == sign(y))`
+1. `MUL` $v_q$, $r_y$, --, $v_{qy}$   `// compute q * y`
+1. `ADD` $v_{qy}$, $v_r$, --, $v_0$   `// compute q * y + r`
+1. `ASSERT_EQ` $v_0$, $x$, --, --
+1. `MOVE` $v_r$, --, --, `rd`
+
+## R1CS constraints
+
+### Ciruict flags
+
+With the M extension we introduce the following circuit flags:
+
+1. `is_virtual`: Is this instruction part of a virtual sequence?
+1. `is_assert`: Is this instruction an `ASSERT_*` instruction?
+1. `do_not_update_pc`: If this instruction is virtual and *not the last one in its sequence*,
+then we should *not* update the PC.
+This is because all instructions in virtual sequences are mapped to the same ELF address.
+
+### Uniform constraints
+
+The following constraints are enforced for every step of the execution trace:
+
+1. If the instruction is a `MUL`, `MULU`, or `MULHU`, the lookup query is the product
+of the two operands `x * y` (field multiplication of two 32-bit values).
+1. If the instruction is a `MOV` or `MOVSIGN`, the lookup query is a single operand `x`
+(read from the first source register `rs1`).
+1. If the instruction is an assert, the lookup output must be true.
+
+### Program counter constraints
+
+Each instruction in the preprocessed [bytecode](./bytecode.md) contains its (compressed)
+memory address as given by the ELF file.
+This is used to compute the expected program counter for each step in the program trace.
+
+If the `do_not_update_pc` flag is set, we constrain the next PC value to be equal to the current one.
+This handles the fact that all instructions in virtual sequences are mapped to the same ELF address.
+
+This also means we need some other mechanism to ensure that virtual sequences are executed in *order* and in *full*.
+If the current instruction is virtual, we can constrain the next instruction in the trace to be the
+next instruction in the bytecode.
+We observe that the virtual sequences used in the M extension don't involve jumps or branches,
+so this should always hold, *except* if we encounter a virtual instruction followed by a padding instruction.
+But that should never happend because an execution trace should always end with some return handling,
+which shouldn't involve a virtual sequence.

book/src/how_it_works.md

@@ -1,5 +1,5 @@
 # How it works
-- [Jolt](./how/jolt.md)
+- [Architecture overview](./how/architecture.md)
     - [Instruction lookups](./how/instruction_lookups.md)
     - [Read-write memory](./how/read_write_memory.md)
     - [R1CS constraints](./how/r1cs_constraints.md)

common/src/rv_trace.rs

@@ -114,6 +114,7 @@ impl From<&RVTraceRow> for [MemoryOp; MEMORY_OPS_PER_INSTRUCTION] {
                 | RV32IM::SLTI
                 | RV32IM::SLTIU
                 | RV32IM::JALR
+                | RV32IM::VIRTUAL_MOVE
                 | RV32IM::VIRTUAL_MOVSIGN => [
                     rs1_read(),
                     MemoryOp::noop_read(),
@@ -222,11 +223,14 @@ pub struct ELFInstruction {
     pub rd: Option<u64>,
     pub imm: Option<u32>,
     /// If this instruction is part of a "virtual sequence" (see Section 6.2 of the
-    /// Jolt paper), then this contains the instruction's index within the sequence.
-    pub virtual_sequence_index: Option<usize>,
+    /// Jolt paper), then this contains the number of virtual instructions after this
+    /// one in the sequence. I.e. if this is the last instruction in the sequence,
+    /// `virtual_sequence_remaining` will be Some(0); if this is the penultimate instruction
+    /// in the sequence, `virtual_sequence_remaining` will be Some(1); etc.
+    pub virtual_sequence_remaining: Option<usize>,
 }
 
-pub const NUM_CIRCUIT_FLAGS: usize = 11;
+pub const NUM_CIRCUIT_FLAGS: usize = 12;
 
 impl ELFInstruction {
     #[rustfmt::skip]
@@ -241,8 +245,9 @@ impl ELFInstruction {
         // 6: Instruction writes lookup output to rd
         // 7: Sign-bit of imm
         // 8: Is concat
-        // 9: Increment virtual PC
+        // 9: Virtual instruction
         // 10: Assert instruction
+        // 11: Don't update PC
 
         let mut flags = [false; NUM_CIRCUIT_FLAGS];
 
@@ -287,7 +292,7 @@ impl ELFInstruction {
             RV32IM::BEQ | RV32IM::BNE | RV32IM::BLT | RV32IM::BGE | RV32IM::BLTU | RV32IM::BGEU,
         );
 
-        // loads, stores, branches, jumps do not store the lookup output to rd (they may update rd in other ways)
+        // loads, stores, branches, jumps, and asserts do not store the lookup output to rd (they may update rd in other ways)
         flags[6] = !matches!(
             self.opcode,
             RV32IM::SB
@@ -301,7 +306,12 @@ impl ELFInstruction {
             | RV32IM::BGEU
             | RV32IM::JAL
             | RV32IM::JALR
-            | RV32IM::LUI,
+            | RV32IM::LUI
+            | RV32IM::VIRTUAL_ASSERT_EQ
+            | RV32IM::VIRTUAL_ASSERT_LTE
+            | RV32IM::VIRTUAL_ASSERT_VALID_DIV0
+            | RV32IM::VIRTUAL_ASSERT_VALID_SIGNED_REMAINDER
+            | RV32IM::VIRTUAL_ASSERT_VALID_UNSIGNED_REMAINDER,
         );
 
         let mask = 1u32 << 31;
@@ -330,23 +340,16 @@ impl ELFInstruction {
             | RV32IM::BLT
             | RV32IM::BGE
             | RV32IM::BLTU
-            | RV32IM::BGEU,
+            | RV32IM::BGEU
+            | RV32IM::VIRTUAL_ASSERT_EQ
+            | RV32IM::VIRTUAL_ASSERT_LTE
+            | RV32IM::VIRTUAL_ASSERT_VALID_SIGNED_REMAINDER
+            | RV32IM::VIRTUAL_ASSERT_VALID_UNSIGNED_REMAINDER
+            | RV32IM::VIRTUAL_ASSERT_VALID_DIV0,
         );
 
-        // TODO(moodlezoup): Use these flags in R1CS constraints
-        flags[9] = match self.virtual_sequence_index {
-            // For virtual sequences, we set
-            //     virtual PC := ProgARW     (the bytecode `a` value)
-            // if it's the first instruction in the sequence.
-            // Otherwise, we increment the virtual PC:
-            //     virtual PC += 1
-            // This prevents a malicious prover from reordering or omitting
-            // instructions from the virtual sequence.
-            Some(i) => i > 0,
-            // For "real" instructions, we always set
-            //     virtual PC := ProgARW     (the bytecode `a` value)
-            None => false
-        };
+        flags[9] = self.virtual_sequence_remaining.is_some();
+
         flags[10] = matches!(self.opcode,
             RV32IM::VIRTUAL_ASSERT_EQ                        |
             RV32IM::VIRTUAL_ASSERT_LTE                       |
@@ -355,6 +358,14 @@ impl ELFInstruction {
             RV32IM::VIRTUAL_ASSERT_VALID_DIV0,
         );
 
+        // All instructions in virtual sequence are mapped from the same
+        // ELF address. Thus if an instruction is virtual (and not the last one
+        // in its sequence), then we should *not* update the PC.
+        flags[11] = match self.virtual_sequence_remaining {
+            Some(i) => i != 0,
+            None => false
+        };
+
         flags
     }
 }
@@ -444,6 +455,7 @@ pub enum RV32IM {
     UNIMPL,
     // Virtual instructions
     VIRTUAL_MOVSIGN,
+    VIRTUAL_MOVE,
     VIRTUAL_ADVICE,
     VIRTUAL_ASSERT_LTE,
     VIRTUAL_ASSERT_VALID_UNSIGNED_REMAINDER,
@@ -546,16 +558,17 @@ impl RV32IM {
             RV32IM::REM    |
             RV32IM::REMU => RV32InstructionFormat::R,
 
-            RV32IM::ADDI  |
-            RV32IM::XORI  |
-            RV32IM::ORI   |
-            RV32IM::ANDI  |
-            RV32IM::SLLI  |
-            RV32IM::SRLI  |
-            RV32IM::SRAI  |
-            RV32IM::SLTI  |
-            RV32IM::FENCE |
-            RV32IM::SLTIU |
+            RV32IM::ADDI         |
+            RV32IM::XORI         |
+            RV32IM::ORI          |
+            RV32IM::ANDI         |
+            RV32IM::SLLI         |
+            RV32IM::SRLI         |
+            RV32IM::SRAI         |
+            RV32IM::SLTI         |
+            RV32IM::FENCE        |
+            RV32IM::SLTIU        |
+            RV32IM::VIRTUAL_MOVE |
             RV32IM::VIRTUAL_MOVSIGN=> RV32InstructionFormat::I,
 
             RV32IM::LB  |

examples/muldiv/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "muldiv"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
+guest = { package = "muldiv-guest", path = "./guest" }
+

examples/muldiv/guest/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "muldiv-guest"
+version = "0.1.0"
+edition = "2021"
+
+[[bin]]
+name = "guest"
+path = "./src/lib.rs"
+
+[features]
+guest = []
+
+[dependencies]
+jolt = { package = "jolt-sdk", path = "../../../jolt-sdk" }

examples/muldiv/guest/src/lib.rs

@@ -0,0 +1,7 @@
+#![cfg_attr(feature = "guest", no_std)]
+#![no_main]
+
+#[jolt::provable]
+fn muldiv(a: u32, b: u32, c: u32) -> u32 {
+    a * b / c
+}

examples/muldiv/src/main.rs

@@ -0,0 +1,9 @@
+pub fn main() {
+    let (prove, verify) = guest::build_muldiv();
+
+    let (output, proof) = prove(12031293, 17, 92);
+    let is_valid = verify(proof);
+
+    println!("output: {}", output);
+    println!("valid: {}", is_valid);
+}

guest-toolchain-tag

@@ -0,0 +1 @@
+nightly-8af9d45d5e09a04832cc9b2e1df993fd1ce49d02