add qm31 operations to VM

ohad-nir-starkware · ohad-nir-starkware · commit 807466154763 · 2025-02-19T14:40:55.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: implement an opcode that computes QM31 arithmetics (add, sub, mul, div) in the VM [#1938](https://github.com/lambdaclass/cairo-vm/pull/1938)
+
 * feat: add functions that compute packed reduced qm31 arithmetics to `math_utils` [#1944](https://github.com/lambdaclass/cairo-vm/pull/1944)
 
 * feat: implement `Blake2sLastBlock` opcode in VM [#1932](https://github.com/lambdaclass/cairo-vm/pull/1932)
diff --git a/cairo_programs/stwo_exclusive_programs/qm31_opcodes_test.cairo b/cairo_programs/stwo_exclusive_programs/qm31_opcodes_test.cairo
@@ -0,0 +1,181 @@
+from starkware.cairo.common.alloc import alloc
+from starkware.cairo.common.bool import FALSE, TRUE
+
+// Tests the QM31_add_mul opcode runners using specific examples as reference.
+// The test is comprised of 10 test cases, each of which tests a different combination of missing_operand, is_imm, and is_mul.
+// is_mul determines whether the operation is a multiplication or an addition of QM31 elements.
+// is_imm determines whether op1 is an immediate.
+// missing_operand determines which operand is missing and needs to be computed by the VM (0 for dst, 1 for op0, 2 fo op1).
+// the combination of is_imm=TRUE with missig_operand=2 is not tested because we do not use arithmetic opcodes to deduce immediates.
+func main{}() {
+    let qm31_op0_coordinates_a = 0x544b2fba;
+    let qm31_op0_coordinates_b = 0x673cff77;
+    let qm31_op0_coordinates_c = 0x60713d44;
+    let qm31_op0_coordinates_d = 0x499602d2;
+    let qm31_op0 = qm31_op0_coordinates_a + qm31_op0_coordinates_b*(2**36) + qm31_op0_coordinates_c*(2**72) + qm31_op0_coordinates_d*(2**108);
+
+    let qm31_op1_coordinates_a = 0x4b18de99;
+    let qm31_op1_coordinates_b = 0x55f6fb62;
+    let qm31_op1_coordinates_c = 0x6e2290d9;
+    let qm31_op1_coordinates_d = 0x7cd851b9;
+    let qm31_op1 = qm31_op1_coordinates_a + qm31_op1_coordinates_b*(2**36) + qm31_op1_coordinates_c*(2**72) + qm31_op1_coordinates_d*(2**108);
+
+    let qm31_add_dst_coordinates_a = 0x1f640e54;
+    let qm31_add_dst_coordinates_b = 0x3d33fada;
+    let qm31_add_dst_coordinates_c = 0x4e93ce1e;
+    let qm31_add_dst_coordinates_d = 0x466e548c;
+    let qm31_add_dst = qm31_add_dst_coordinates_a + qm31_add_dst_coordinates_b*(2**36) + qm31_add_dst_coordinates_c*(2**72) + qm31_add_dst_coordinates_d*(2**108);
+
+    let qm31_mul_dst_coordinates_a = 0x38810ab4;
+    let qm31_mul_dst_coordinates_b = 0x5a0fd30a;
+    let qm31_mul_dst_coordinates_c = 0x2527b81e;
+    let qm31_mul_dst_coordinates_d = 0x4b1ed1cd;
+    let qm31_mul_dst = qm31_mul_dst_coordinates_a + qm31_mul_dst_coordinates_b*(2**36) + qm31_mul_dst_coordinates_c*(2**72) + qm31_mul_dst_coordinates_d*(2**108);
+
+    let runner_output_mul_dst = run_qm31_operation(missing_operand=0, is_imm=FALSE, is_mul=TRUE, dst_or_op0=qm31_op0, op0_or_op1=qm31_op1);
+    assert runner_output_mul_dst = qm31_mul_dst;
+    let runner_output_add_dst = run_qm31_operation(missing_operand=0, is_imm=FALSE, is_mul=FALSE, dst_or_op0=qm31_op0, op0_or_op1=qm31_op1);
+    assert runner_output_add_dst = qm31_add_dst;
+
+    let runner_output_mul_op0 = run_qm31_operation(missing_operand=1, is_imm=FALSE, is_mul=TRUE, dst_or_op0=qm31_mul_dst, op0_or_op1=qm31_op1);
+    assert runner_output_mul_op0 = qm31_op0;
+    let runner_output_add_op0 = run_qm31_operation(missing_operand=1, is_imm=FALSE, is_mul=FALSE, dst_or_op0=qm31_add_dst, op0_or_op1=qm31_op1);
+    assert runner_output_add_op0 = qm31_op0;
+
+    let runner_output_mul_op1 = run_qm31_operation(missing_operand=2, is_imm=FALSE, is_mul=TRUE, dst_or_op0=qm31_mul_dst, op0_or_op1=qm31_op0);
+    assert runner_output_mul_op1 = qm31_op1;
+    let runner_output_add_op1 = run_qm31_operation(missing_operand=2, is_imm=FALSE, is_mul=FALSE, dst_or_op0=qm31_add_dst, op0_or_op1=qm31_op0);
+    assert runner_output_add_op1 = qm31_op1;
+
+    let runner_output_mul_dst = run_qm31_operation(missing_operand=0, is_imm=TRUE, is_mul=TRUE, dst_or_op0=qm31_op0, op0_or_op1=qm31_op1);
+    assert runner_output_mul_dst = qm31_mul_dst;
+    let runner_output_add_dst = run_qm31_operation(missing_operand=0, is_imm=TRUE, is_mul=FALSE, dst_or_op0=qm31_op0, op0_or_op1=qm31_op1);
+    assert runner_output_add_dst = qm31_add_dst;
+
+    let runner_output_mul_op0 = run_qm31_operation(missing_operand=1, is_imm=TRUE, is_mul=TRUE, dst_or_op0=qm31_mul_dst, op0_or_op1=qm31_op1);
+    assert runner_output_mul_op0 = qm31_op0;
+    let runner_output_add_op0 = run_qm31_operation(missing_operand=1, is_imm=TRUE, is_mul=FALSE, dst_or_op0=qm31_add_dst, op0_or_op1=qm31_op1);
+    assert runner_output_add_op0 = qm31_op0;
+
+    return ();
+}
+
+// Forces the runner to execute the QM31_add_mul opcode with the given operands.
+// missing_operand, is_imm, is_mul determine the configuration of the operation as described above.
+// dst_or_op0 is a felt representing the value of either the op0 (if missing_operand=0) or dst (otherwise) operand.
+// op0_or_op1 is a felt representing the value of either the op0 (if missing_operand=2) or op1 (otherwise) operand.
+// dst_or_op0 and op0_or_op1 are stored within addresses fp-4 and fp-3 respectively, they are passed to the instruction
+// using offsets wrt fp (unless is_imm=TRUE, in which case op1 has offset 1 relative to pc).
+// The missing operand has offset 0 relative to ap.
+// An instruction encoding with the appropriate flags and offsets is built, then written to [pc] and the runner is forced to execute QM31_add_mul.
+// The missing operand is deduced to [ap] and returned.
+func run_qm31_operation(
+    missing_operand: felt,
+    is_imm: felt,
+    is_mul: felt,
+    dst_or_op0: felt,
+    op0_or_op1: felt,
+) -> felt {
+    alloc_locals;
+    
+    // Set flags and offsets.
+    let (local offsets) = alloc();
+    let (local flags) = alloc();
+
+    assert offsets[missing_operand] = 2**15; // the missing operand will be written to [ap]
+
+    assert flags[2] = is_imm; // flag_op1_imm = 0;
+    assert flags[5] = 1-is_mul; // flag_res_add = 1-is_mul;
+    assert flags[6] = is_mul; // flag_res_mul = is_mul;
+    assert flags[7] = 0; // flag_PC_update_jump = 0;
+    assert flags[8] = 0; // flag_PC_update_jump_rel = 0;
+    assert flags[9] = 0; // flag_PC_update_jnz = 0;
+    assert flags[10] = 0; // flag_ap_update_add = 0;
+    assert flags[11] = 0; // flag_ap_update_add_1 = 0;
+    assert flags[12] = 0; // flag_opcode_call = 0;
+    assert flags[13] = 0; // flag_opcode_ret = 0;
+    assert flags[14] = 1; // flag_opcode_assert_eq = 1;
+
+    if (missing_operand == 0) {
+        assert offsets[1] = 2**15 - 4;
+        assert offsets[2] = 2**15 - 3 + 4 * is_imm;
+        assert flags[0] = 0; // flag_dst_base_fp
+        assert flags[1] = 1; // flag_op0_base_fp
+    }
+    if (missing_operand == 1) {
+        assert offsets[0] = 2**15 - 4;
+        assert offsets[2] = 2**15 - 3 + 4 * is_imm;
+        assert flags[0] = 1; // flag_dst_base_fp
+        assert flags[1] = 0; // flag_op0_base_fp
+    }
+    if (missing_operand == 2) {
+        assert is_imm = FALSE;
+        assert offsets[0] = 2**15 - 4;
+        assert offsets[1] = 2**15 - 3;
+        assert flags[0] = 1; // flag_dst_base_fp
+        assert flags[1] = 1; // flag_op0_base_fp
+    }
+    assert flags[3] = 2 - is_imm - flags[0] - flags[1]; // flag_op1_base_fp
+    assert flags[4] = 1 - is_imm - flags[3]; // flag_op1_base_ap
+
+    // Compute the instruction encoding.
+    let flag_num = flags[0] + flags[1]*(2**1) + flags[2]*(2**2) + flags[3]*(2**3) + flags[4]*(2**4) + flags[5]*(2**5) + flags[6]*(2**6) + flags[14]*(2**14);
+    let qm31_opcode_extension_num = 3;
+    let instruction_encoding = offsets[0] + offsets[1]*(2**16) + offsets[2]*(2**32) + flag_num*(2**48) + qm31_opcode_extension_num*(2**63);
+
+    // Run the instruction and return the result.
+    if (is_imm == TRUE) {
+        assert op0_or_op1 = 0x7cd851b906e2290d9055f6fb6204b18de99;
+        if (missing_operand == 0) {
+            if (is_mul == TRUE) {
+                assert instruction_encoding=0x1c04680017ffc8000;
+                dw 0x1c04680017ffc8000;
+                dw 0x7cd851b906e2290d9055f6fb6204b18de99;
+                return [ap];
+            }
+            assert instruction_encoding=0x1c02680017ffc8000;
+            dw 0x1c02680017ffc8000;
+            dw 0x7cd851b906e2290d9055f6fb6204b18de99;
+            return [ap];
+        }
+        if (missing_operand == 1) {
+            if (is_mul == TRUE) {
+                assert instruction_encoding=0x1c045800180007ffc;
+                dw 0x1c045800180007ffc;
+                dw 0x7cd851b906e2290d9055f6fb6204b18de99;
+                return [ap];
+            }
+            assert instruction_encoding=0x1c025800180007ffc;
+            dw 0x1c025800180007ffc;
+            dw 0x7cd851b906e2290d9055f6fb6204b18de99;
+            return [ap];
+        }
+    }
+
+    if (missing_operand == 0) {
+        if (is_mul == TRUE) {
+            assert instruction_encoding=0x1c04a7ffd7ffc8000;
+            dw 0x1c04a7ffd7ffc8000;
+            return [ap];
+        }
+        assert instruction_encoding=0x1c02a7ffd7ffc8000;
+        dw 0x1c02a7ffd7ffc8000;
+        return [ap];
+    }
+    if (missing_operand == 1) {
+        if (is_mul == TRUE) {
+            assert instruction_encoding=0x1c0497ffd80007ffc;
+            dw 0x1c0497ffd80007ffc;
+            return [ap];
+        }
+        assert instruction_encoding=0x1c0297ffd80007ffc;
+        dw 0x1c0297ffd80007ffc;
+        return [ap];
+    }
+    if (is_mul == TRUE) {
+        dw 0x1c05380007ffd7ffc;
+        return [ap];
+    }
+    dw 0x1c03380007ffd7ffc;
+    return [ap];
+}
diff --git a/vm/src/tests/cairo_run_test.rs b/vm/src/tests/cairo_run_test.rs
@@ -576,6 +576,14 @@ fn blake2s_opcode_test() {
     run_program_simple(program_data.as_slice());
 }
 
+#[test]
+#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+fn qm31_opcodes_test() {
+    let program_data =
+        include_bytes!("../../../cairo_programs/stwo_exclusive_programs/qm31_opcodes_test.json");
+    run_program_simple(program_data.as_slice());
+}
+
 #[test]
 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
 fn relocate_segments() {
diff --git a/vm/src/types/errors/math_errors.rs b/vm/src/types/errors/math_errors.rs
@@ -42,6 +42,24 @@ pub enum MathError {
     RelocatableAddUsizeOffsetExceeded(Box<(Relocatable, usize)>),
     #[error("Operation failed: {} + {}, can't add two relocatable values", (*.0).0, (*.0).1)]
     RelocatableAdd(Box<(Relocatable, Relocatable)>),
+    #[error("Operation failed: {} - {}, can't add a relocatable value as a QM31 element", (*.0).0, (*.0).1)]
+    RelocatableQM31Add(Box<(MaybeRelocatable, MaybeRelocatable)>),
+    #[error("Operation failed: {} - {}, can't subtract a relocatable value or from a relocatable value as a QM31 element", (*.0).0, (*.0).1)]
+    RelocatableQM31Sub(Box<(MaybeRelocatable, MaybeRelocatable)>),
+    #[error("Operation failed: {} * {}, can't multiply a relocatable value", (*.0).0, (*.0).1)]
+    RelocatableMul(Box<(MaybeRelocatable, MaybeRelocatable)>),
+    #[error(
+        "Operation failed: typed_add for an OpcodeExtension that is neither Stone nor QM31Operation"
+    )]
+    InvalidAddOpcodeExtension(),
+    #[error(
+        "Operation failed: typed_mul for an OpcodeExtension that is neither Stone nor QM31Operation"
+    )]
+    InvalidSubOpcodeExtension(),
+    #[error(
+        "Operation failed: typed_sub for an OpcodeExtension that is neither Stone nor QM31Operation"
+    )]
+    InvalidMulOpcodeExtension(),
     #[error("Operation failed: {} - {}, can't subtract two relocatable values with different segment indexes", (*.0).0, (*.0).1)]
     RelocatableSubDiffIndex(Box<(Relocatable, Relocatable)>),
     #[error(
diff --git a/vm/src/types/instruction.rs b/vm/src/types/instruction.rs
@@ -82,6 +82,7 @@ pub enum OpcodeExtension {
     Stone,
     Blake,
     BlakeFinalize,
+    QM31Operation,
 }
 
 impl Instruction {
diff --git a/vm/src/types/relocatable.rs b/vm/src/types/relocatable.rs
@@ -4,9 +4,14 @@ use crate::stdlib::{
     prelude::*,
 };
 
+use crate::math_utils::{
+    qm31_packed_reduced_add, qm31_packed_reduced_mul, qm31_packed_reduced_sub,
+};
 use crate::Felt252;
 use crate::{
-    relocatable, types::errors::math_errors::MathError, vm::errors::memory_errors::MemoryError,
+    relocatable,
+    types::{errors::math_errors::MathError, instruction::OpcodeExtension},
+    vm::errors::memory_errors::MemoryError,
 };
 use num_traits::ToPrimitive;
 use serde::{Deserialize, Serialize};
@@ -330,6 +335,89 @@ impl MaybeRelocatable {
         }
     }
 
+    /// Substracts a MaybeRelocatable value from self according to the specified OpcodeExtension.
+    /// If the OpcodeExtension is Stone it subtracts as MaybeRelocatable::sub does.
+    /// If the OpcodeExtension is QM31Operation it requires them both to be Int and it subtracts
+    /// them as packed reduced QM31 elements.
+    pub fn typed_add(
+        &self,
+        other: &MaybeRelocatable,
+        opcode_extension: OpcodeExtension,
+    ) -> Result<MaybeRelocatable, MathError> {
+        match opcode_extension {
+            OpcodeExtension::Stone => self.add(other),
+            OpcodeExtension::QM31Operation => {
+                if let (MaybeRelocatable::Int(num_self), MaybeRelocatable::Int(num_other)) =
+                    (self, other)
+                {
+                    Ok(MaybeRelocatable::Int(qm31_packed_reduced_add(
+                        *num_self, *num_other,
+                    )?))
+                } else {
+                    Err(MathError::RelocatableQM31Add(Box::new((
+                        self.clone(),
+                        other.clone(),
+                    ))))
+                }
+            }
+            _ => Err(MathError::InvalidAddOpcodeExtension()),
+        }
+    }
+
+    /// Substracts a MaybeRelocatable value from self according to the specified OpcodeExtension.
+    /// If the OpcodeExtension is Stone it subtracts as MaybeRelocatable::sub does.
+    /// If the OpcodeExtension is QM31Operation it requires them both to be Int and it subtracts
+    /// them as packed reduced QM31 elements.
+    pub fn typed_sub(
+        &self,
+        other: &MaybeRelocatable,
+        opcode_extension: OpcodeExtension,
+    ) -> Result<MaybeRelocatable, MathError> {
+        match opcode_extension {
+            OpcodeExtension::Stone => self.sub(other),
+            OpcodeExtension::QM31Operation => {
+                if let (MaybeRelocatable::Int(num_self), MaybeRelocatable::Int(num_other)) =
+                    (self, other)
+                {
+                    Ok(MaybeRelocatable::Int(qm31_packed_reduced_sub(
+                        *num_self, *num_other,
+                    )?))
+                } else {
+                    Err(MathError::RelocatableQM31Sub(Box::new((
+                        self.clone(),
+                        other.clone(),
+                    ))))
+                }
+            }
+            _ => Err(MathError::InvalidSubOpcodeExtension()),
+        }
+    }
+
+    /// Multiplies self and another MaybeRelocatable value according to the specified OpcodeExtension.
+    /// Requires both operands to be Int.
+    /// If the OpcodeExtension is Stone it multiplies them as Felts.
+    /// If the OpcodeExtension is QM31Operation it multiplies them as packed reduced QM31 elements.
+    pub fn typed_mul(
+        &self,
+        other: &MaybeRelocatable,
+        opcode_extension: OpcodeExtension,
+    ) -> Result<MaybeRelocatable, MathError> {
+        if let (MaybeRelocatable::Int(num_self), MaybeRelocatable::Int(num_other)) = (self, other) {
+            match opcode_extension {
+                OpcodeExtension::Stone => Ok(MaybeRelocatable::Int(num_self * num_other)),
+                OpcodeExtension::QM31Operation => Ok(MaybeRelocatable::Int(
+                    qm31_packed_reduced_mul(*num_self, *num_other)?,
+                )),
+                _ => Err(MathError::InvalidMulOpcodeExtension()),
+            }
+        } else {
+            Err(MathError::RelocatableMul(Box::new((
+                self.clone(),
+                other.clone(),
+            ))))
+        }
+    }
+
     // TODO: Check if its more performant to use get_int instead
     /// Returns a reference to the inner value if it is a Felt252, returns None otherwise.
     pub fn get_int_ref(&self) -> Option<&Felt252> {
@@ -678,6 +766,48 @@ mod tests {
         );
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn typed_add_blake() {
+        let a = &MaybeRelocatable::from(5);
+        let b = &MaybeRelocatable::from(6);
+        let error = a.typed_add(b, OpcodeExtension::Blake);
+        assert_eq!(error, Err(MathError::InvalidAddOpcodeExtension()));
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn typed_sub_blake() {
+        let a = &MaybeRelocatable::from(7);
+        let b = &MaybeRelocatable::from(3);
+        let error = a.typed_sub(b, OpcodeExtension::Blake);
+        assert_eq!(error, Err(MathError::InvalidSubOpcodeExtension()));
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn relocatable_typed_sub_q31_operation() {
+        let a = &MaybeRelocatable::from((6, 8));
+        let b = &MaybeRelocatable::from(2);
+        let error = a.typed_sub(b, OpcodeExtension::QM31Operation);
+        assert_eq!(
+            error,
+            Err(MathError::RelocatableQM31Sub(Box::new((
+                MaybeRelocatable::from((6, 8)),
+                MaybeRelocatable::from(2)
+            ))))
+        );
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn typed_mul_blake_finalize() {
+        let a = &MaybeRelocatable::from(4);
+        let b = &MaybeRelocatable::from(9);
+        let error = a.typed_mul(b, OpcodeExtension::BlakeFinalize);
+        assert_eq!(error, Err(MathError::InvalidMulOpcodeExtension()));
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn divmod_working() {
diff --git a/vm/src/vm/decoding/decoder.rs b/vm/src/vm/decoding/decoder.rs
diff --git a/vm/src/vm/errors/vm_errors.rs b/vm/src/vm/errors/vm_errors.rs
diff --git a/vm/src/vm/vm_core.rs b/vm/src/vm/vm_core.rs

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ pub enum OpcodeExtension {`
`82`	`82`	`Stone,`
`83`	`83`	`Blake,`
`84`	`84`	`BlakeFinalize,`
	`85`	`+ QM31Operation,`
`85`	`86`	`}`
`86`	`87`
`87`	`88`	`impl Instruction {`