Feat/MULH opcode

ceno_emul/src/rv32im_encode.rs

@@ -8,6 +8,13 @@ const MASK_8_BITS: u32 = 0xFF;
 const MASK_10_BITS: u32 = 0x3FF;
 const MASK_12_BITS: u32 = 0xFFF;
 
+/// Generate bit encoding of a RISC-V instruction.
+///
+/// Values `rs1`, `rs2` and `rd1` are 5-bit register indices, and `imm` is of
+/// bit length depending on the requirements of the instruction format type.
+///
+/// Fields not required by the instruction's format type are ignored, so one can
+/// safely pass an arbitrary value for these, say 0.
 pub const fn encode_rv32(kind: InsnKind, rs1: u32, rs2: u32, rd: u32, imm: u32) -> u32 {
     match kind.codes().format {
         InsnFormat::R => encode_r(kind, rs1, rs2, rd),

ceno_zkvm/src/instructions/riscv/mulh.rs

@@ -1,15 +1,24 @@
-use std::marker::PhantomData;
+use std::{fmt::Display, marker::PhantomData};
 
 use ceno_emul::{InsnKind, StepRecord};
 use ff_ext::ExtensionField;
 
-use super::{
-    RIVInstruction,
-    constants::{UInt, UIntMul},
-    r_insn::RInstructionConfig,
-};
 use crate::{
-    circuit_builder::CircuitBuilder, error::ZKVMError, instructions::Instruction, uint::Value,
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{Expression, ToExpr, WitIn},
+    gadgets::IsLtConfig,
+    instructions::{
+        Instruction,
+        riscv::{
+            RIVInstruction,
+            constants::{BIT_WIDTH, LIMB_BITS, UInt, UIntMul},
+            r_insn::RInstructionConfig,
+        },
+    },
+    set_val,
+    uint::Value,
+    utils::i64_to_base,
     witness::LkMultiplicity,
 };
 use core::mem::MaybeUninit;
@@ -24,15 +33,15 @@ pub struct ArithConfig<E: ExtensionField> {
     rd_written: UIntMul<E>,
 }
 
-pub struct MulhInstruction<E, I>(PhantomData<(E, I)>);
+pub struct MulhInstructionBase<E, I>(PhantomData<(E, I)>);
 
 pub struct MulhuOp;
 impl RIVInstruction for MulhuOp {
     const INST_KIND: InsnKind = InsnKind::MULHU;
 }
-pub type MulhuInstruction<E> = MulhInstruction<E, MulhuOp>;
+pub type MulhuInstruction<E> = MulhInstructionBase<E, MulhuOp>;
 
-impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for MulhInstruction<E, I> {
+impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for MulhInstructionBase<E, I> {
     type InstructionConfig = ArithConfig<E>;
 
     fn name() -> String {
@@ -115,6 +124,213 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for MulhInstruction<E,
     }
 }
 
+pub struct MulhInstruction<E>(PhantomData<E>);
+
+pub struct MulhConfig<E: ExtensionField> {
+    rs1_read: UInt<E>,
+    rs2_read: UInt<E>,
+    rd_written: UInt<E>,
+    rs1_signed: Signed,
+    rs2_signed: Signed,
+    rd_sign_bit: IsLtConfig,
+    unsigned_prod_low: UInt<E>,
+    r_insn: RInstructionConfig<E>,
+}
+
+impl<E: ExtensionField> Instruction<E> for MulhInstruction<E> {
+    type InstructionConfig = MulhConfig<E>;
+
+    fn name() -> String {
+        format!("{:?}", InsnKind::MULH)
+    }
+
+    /// Circuit is validated by the following strategy:
+    /// 1. Compute the signed values associated with `rs1` and `rs2`
+    /// 2. Compute the high order bit of `rd`, which is the sign bit of the 2s
+    ///    complement value for which rd represents the high limb
+    /// 3. Verify that the product of signed inputs `rs1` and `rs2` is equal to
+    ///    the result of interpreting rd as the high limb of a 2s complement
+    ///    value with some 32-bit low limb
+    ///
+    /// The correctness here is a bit subtle.  The signed values of 32-bit
+    /// inputs `rs1` and `rs2` have values between `-2^31` and `2^31 - 1`, so
+    /// their product is constrained to lie between `-2^62 + 2^31` and
+    /// `2^62`.  In a prime field of size smaller than `2^64`, the range of
+    /// values represented by a 64-bit 2s complement value, integers between
+    /// `-2^63` and `2^63 - 1`, have some ambiguity.  If `p = 2^64 - k`, then
+    /// the values between `-2^63` and `-2^63 + k - 1` correspond with the values
+    /// between `2^63 - k` and `2^63 - 1`.
+    ///
+    /// However, as long as the values required by signed products don't overlap
+    /// with this ambiguous range, an arbitrary 64-bit 2s complement value can
+    /// represent a signed 32-bit product in only one way, so there is no
+    /// ambiguity in the representation.  This is the case for the Goldilocks
+    /// field with order `p = 2^64 - 2^32 + 1`.
+    fn construct_circuit(
+        circuit_builder: &mut CircuitBuilder<E>,
+    ) -> Result<MulhConfig<E>, ZKVMError> {
+        let rs1_read = UInt::new_unchecked(|| "rs1_read", circuit_builder)?;
+        let rs2_read = UInt::new_unchecked(|| "rs2_read", circuit_builder)?;
+        let rd_written = UInt::new(|| "rd_written", circuit_builder)?;
+
+        let rs1_signed = Signed::construct_circuit(circuit_builder, || "rs1", &rs1_read)?;
+        let rs2_signed = Signed::construct_circuit(circuit_builder, || "rs2", &rs2_read)?;
+
+        let rd_sign_bit = IsLtConfig::construct_circuit(
+            circuit_builder,
+            || "rd_sign_bit",
+            ((1u64 << (LIMB_BITS - 1)) - 1).into(),
+            rd_written.expr().last().unwrap().clone(),
+            1,
+        )?;
+
+        let unsigned_prod_low = UInt::new(|| "unsigned_prod_low", circuit_builder)?;
+        circuit_builder.require_equal(
+            || "validate_prod_high_limb",
+            rs1_signed.val.expr() * rs2_signed.val.expr(),
+            Expression::<E>::from(1u64 << 32) * rd_written.value() + unsigned_prod_low.value()
+                - Expression::<E>::from(1u128 << 64) * rd_sign_bit.expr(),
+        )?;
+
+        let r_insn = RInstructionConfig::<E>::construct_circuit(
+            circuit_builder,
+            InsnKind::MULH,
+            rs1_read.register_expr(),
+            rs2_read.register_expr(),
+            rd_written.register_expr(),
+        )?;
+
+        Ok(MulhConfig {
+            rs1_read,
+            rs2_read,
+            rd_written,
+            rs1_signed,
+            rs2_signed,
+            rd_sign_bit,
+            unsigned_prod_low,
+            r_insn,
+        })
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [MaybeUninit<<E as ExtensionField>::BaseField>],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        // Read registers from step
+        let rs1_read = Value::new_unchecked(step.rs1().unwrap().value);
+        config
+            .rs1_read
+            .assign_limbs(instance, rs1_read.as_u16_limbs());
+
+        let rs2_read = Value::new_unchecked(step.rs2().unwrap().value);
+        config
+            .rs2_read
+            .assign_limbs(instance, rs2_read.as_u16_limbs());
+
+        let rd_written = Value::new(step.rd().unwrap().value.after, lk_multiplicity);
+        config
+            .rd_written
+            .assign_limbs(instance, rd_written.as_u16_limbs());
+
+        // Signed register values
+        let rs1_signed =
+            config
+                .rs1_signed
+                .assign_instance::<E>(instance, lk_multiplicity, &rs1_read)?;
+
+        let rs2_signed =
+            config
+                .rs2_signed
+                .assign_instance::<E>(instance, lk_multiplicity, &rs2_read)?;
+
+        // Sign bit of rd register
+        let rd_high_limb = *rd_written.limbs.last().unwrap() as u64;
+        let sign_cutoff = (1u64 << (LIMB_BITS - 1)) - 1;
+        config
+            .rd_sign_bit
+            .assign_instance(instance, lk_multiplicity, sign_cutoff, rd_high_limb)?;
+
+        // Low limb of product in 2s complement form
+        let prod = ((rs1_signed as i64) * (rs2_signed as i64)) as u64;
+        let unsigned_prod_low = (prod % (1u64 << BIT_WIDTH)) as u32;
+        let unsigned_prod_low_val = Value::new(unsigned_prod_low, lk_multiplicity);
+        config
+            .unsigned_prod_low
+            .assign_limbs(instance, unsigned_prod_low_val.as_u16_limbs());
+
+        // R-type instruction
+        config
+            .r_insn
+            .assign_instance(instance, lk_multiplicity, step)?;
+
+        Ok(())
+    }
+}
+
+/// Transform a value represented as a `UInt` into a `WitIn` containing its
+/// corresponding signed value, interpreting the bits as a 2s-complement
+/// encoding.  Gadget allocates 3 `WitIn` values in total.
+struct Signed {
+    pub is_negative: IsLtConfig,
+    pub val: WitIn,
+}
+
+impl Signed {
+    pub fn construct_circuit<
+        E: ExtensionField,
+        NR: Into<String> + Display + Clone,
+        N: FnOnce() -> NR,
+    >(
+        cb: &mut CircuitBuilder<E>,
+        name_fn: N,
+        unsigned_val: &UInt<E>,
+    ) -> Result<Self, ZKVMError> {
+        cb.namespace(
+            || "signed",
+            |cb| {
+                let name = name_fn();
+                // is_lt is set if top limb of val is negative
+                let is_negative = IsLtConfig::construct_circuit(
+                    cb,
+                    || name.clone(),
+                    ((1u64 << (LIMB_BITS - 1)) - 1).into(),
+                    unsigned_val.expr().last().unwrap().clone(),
+                    1,
+                )?;
+                let val = cb.create_witin(|| format!("{name} signed_val witin"));
+                cb.require_equal(
+                    || "signed_val",
+                    val.expr(),
+                    unsigned_val.value() - (1u64 << BIT_WIDTH) * is_negative.expr(),
+                )?;
+
+                Ok(Self { is_negative, val })
+            },
+        )
+    }
+
+    pub fn assign_instance<E: ExtensionField>(
+        &self,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        lkm: &mut LkMultiplicity,
+        val: &Value<u32>,
+    ) -> Result<i32, ZKVMError> {
+        let high_limb = *val.limbs.last().unwrap() as u64;
+        let sign_cutoff = (1u64 << (LIMB_BITS - 1)) - 1;
+        self.is_negative
+            .assign_instance(instance, lkm, sign_cutoff, high_limb)?;
+
+        let signed_val = val.as_u32() as i32;
+
+        let field_elt: E::BaseField = i64_to_base(signed_val as i64);
+        set_val!(instance, self.val, field_elt);
+
+        Ok(signed_val)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use ceno_emul::{Change, StepRecord, encode_rv32};
@@ -130,12 +346,12 @@ mod test {
 
     #[test]
     fn test_opcode_mulhu() {
-        verify(2, 11);
-        verify(u32::MAX, u32::MAX);
-        verify(u16::MAX as u32, u16::MAX as u32);
+        verify_mulhu(2, 11);
+        verify_mulhu(u32::MAX, u32::MAX);
+        verify_mulhu(u16::MAX as u32, u16::MAX as u32);
     }
 
-    fn verify(rs1: u32, rs2: u32) {
+    fn verify_mulhu(rs1: u32, rs2: u32) {
         let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
         let mut cb = CircuitBuilder::new(&mut cs);
         let config = cb
@@ -175,4 +391,64 @@ mod test {
 
         MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
     }
+
+    #[test]
+    fn test_opcode_mulh() {
+        let test_cases = vec![
+            (2, 11),
+            (7, 0),
+            (0, 5),
+            (0, -3),
+            (-19, 0),
+            (0, 0),
+            (-12, -31),
+            (2, -1),
+            (1, i32::MIN),
+            (i32::MAX, -1),
+            (i32::MAX, i32::MIN),
+            (i32::MAX, i32::MAX),
+            (i32::MIN, i32::MIN),
+        ];
+        test_cases
+            .into_iter()
+            .for_each(|(rs1, rs2)| verify_mulh(rs1, rs2));
+    }
+
+    fn verify_mulh(rs1: i32, rs2: i32) {
+        let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
+        let mut cb = CircuitBuilder::new(&mut cs);
+        let config = cb
+            .namespace(|| "mulh", |cb| Ok(MulhInstruction::construct_circuit(cb)))
+            .unwrap()
+            .unwrap();
+
+        let signed_prod_high = ((rs1 as i64).wrapping_mul(rs2 as i64) >> 32) as u32;
+
+        // values assignment
+        let insn_code = encode_rv32(InsnKind::MULH, 2, 3, 4, 0);
+        let (raw_witin, lkm) =
+            MulhInstruction::assign_instances(&config, cb.cs.num_witin as usize, vec![
+                StepRecord::new_r_instruction(
+                    3,
+                    MOCK_PC_START,
+                    insn_code,
+                    rs1 as u32,
+                    rs2 as u32,
+                    Change::new(0, signed_prod_high),
+                    0,
+                ),
+            ])
+            .unwrap();
+
+        // verify value written to register
+        let rd_written_expr = cb.get_debug_expr(DebugIndex::RdWrite as usize)[0].clone();
+        cb.require_equal(
+            || "assert_rd_written",
+            rd_written_expr,
+            Expression::from(signed_prod_high),
+        )
+        .unwrap();
+
+        MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm));
+    }
 }