Build broadcast qm31.

Author: leo-starkware

crates/circuits/src/finalize_constants.rs

@@ -88,6 +88,21 @@ fn finalize_constants_with_min_base(context: &mut Context<impl IValue>, min_base
     // Decompose M31 constants not in the chain by expressing them in base `m31_base`.
     decompose_m31_constants(context, &mut m31_constants, &mut m31_cache, m31_base);
     assert!(m31_constants.is_empty());
+
+    // Deal with the QM31 constants.
+    // Build `i` and `i * u` to get the qm31_basis [i, u, iu].
+    let two = Var { idx: *m31_cache.get(&2.into()).unwrap() };
+    let i_var = eval!(context, ((u_var) * (u_var)) - (two));
+    let iu_var = eval!(context, (i_var) * (u_var));
+    let qm31_basis: [Var; 3] = [i_var, u_var, iu_var];
+    // Build the broadcast QM31 constants, i.e. constants of the form (x, x, x, x), x != 0.
+    decompose_broadcast_constants(
+        context,
+        &mut qm31_constants,
+        &mut m31_cache,
+        m31_base,
+        qm31_basis,
+    );
 }
 
 /// Finds the largest integer N such that all values in [0, N] are present as constants.
@@ -208,3 +223,32 @@ fn build_m31_from_base(
     assert!(m31_cache.contains_key(&val));
     assert!(!m31_constants.contains_key(&val));
 }
+
+fn decompose_broadcast_constants(
+    context: &mut Context<impl IValue>,
+    qm31_constants: &mut IndexMap<QM31, Var>,
+    m31_cache: &mut HashMap<M31, usize>,
+    base: M31,
+    qm31_basis: [Var; 3],
+) {
+    let [i_var, u_var, iu_var] = qm31_basis;
+    let one = context.one();
+    let ones = eval!(context, ((one) + (i_var)) + ((u_var) + (iu_var)));
+
+    qm31_constants.retain(|qm31_value, qm31_var| {
+        let is_broadcast = qm31_value.to_m31_array().iter().tuple_windows().all(|(x, y)| x == y);
+        if !is_broadcast {
+            return true;
+        }
+        let m31_value = qm31_value.0.0;
+        // If m31_value is not in the cache, add it.
+        if !m31_cache.contains_key(&m31_value) {
+            build_m31_from_base(context, m31_cache, &mut IndexMap::new(), base, m31_value);
+        }
+        let m31_idx = *m31_cache.get(&m31_value).unwrap();
+        // Add a gate m31_val * (1, 1, 1, 1) = qm31_var.
+        context.circuit.mul.push(Mul { in0: m31_idx, in1: ones.idx, out: qm31_var.idx });
+        // Remove the element from qm31_constants.
+        false
+    });
+}

crates/circuits/src/finalize_constants_test.rs

@@ -42,7 +42,13 @@ fn test_plus_one_chain_topology() {
         [4] = [6] + [1]
         [7] = [4] + [1]
         [8] = [7] + [1]
+        [12] = [1] + [10]
+        [13] = [2] + [11]
+        [14] = [12] + [13]
+        [10] = [9] - [3]
         [5] = [2] * [1]
+        [9] = [2] * [2]
+        [11] = [10] * [2]
         [5] = [2]
         output [2]
     "#]]
@@ -78,8 +84,14 @@ fn test_large_m31_decomposition() {
         [8] = [7] + [1]
         [9] = [8] + [5]
         [3] = [10] + [5]
+        [14] = [1] + [12]
+        [15] = [2] + [13]
+        [16] = [14] + [15]
+        [12] = [11] - [5]
         [4] = [2] * [1]
         [10] = [9] * [8]
+        [11] = [2] * [2]
+        [13] = [12] * [2]
         [4] = [2]
         output [2]
     "#]]
@@ -92,3 +104,50 @@ fn test_large_m31_decomposition() {
     context.circuit.check_yields();
     context.validate_circuit();
 }
+
+#[test]
+fn test_broadcast_decomposition() {
+    let mut context = TraceContext::default();
+    // Add `u`.
+    // TODO(Leo): remove this once `u` is added to the default constants.
+    context.constant(qm31_from_u32s(0, 0, 1, 0));
+    // Broadcast constant (11, 11, 11, 11) — should be yielded as 11 * (1, 1, 1, 1). Since 11 is
+    // outside the chain (`min_base = 5`), the M31 factor 11 is itself built via base
+    // decomposition: 11 = 2 * 5 + 1.
+    context.constant(qm31_from_u32s(11, 11, 11, 11));
+    finalize_constants_with_min_base(&mut context, 5);
+
+    // The plus-one chain populates [5]..=[8] for values 2..=5. The QM31 basis allocates [9] = u*u,
+    // [10] = u² - 2 = i, [11] = i*u = iu. The ones vector is built as ([1] + [10]) + ([2] + [11])
+    // yielding wires [12], [13], [14]. Then the M31 factor 11 is decomposed in base 5:
+    // [15] = [5] * [8] (= 2 * 5 = 10) and [16] = [15] + [1] (= 11). Finally the broadcast is
+    // yielded by [3] = [16] * [14] (11 * ones).
+    expect![[r#"
+        [0] = [0] + [0]
+        [2] = [2] + [0]
+        [1] = [1] + [0]
+        [5] = [1] + [1]
+        [6] = [5] + [1]
+        [7] = [6] + [1]
+        [8] = [7] + [1]
+        [12] = [1] + [10]
+        [13] = [2] + [11]
+        [14] = [12] + [13]
+        [16] = [15] + [1]
+        [10] = [9] - [5]
+        [4] = [2] * [1]
+        [9] = [2] * [2]
+        [11] = [10] * [2]
+        [15] = [5] * [8]
+        [3] = [16] * [14]
+        [4] = [2]
+        output [2]
+    "#]]
+    .assert_eq(&format!("{:?}", context.circuit));
+
+    assert_eq!(context.get(Var { idx: 14 }), qm31_from_u32s(1, 1, 1, 1));
+    assert_eq!(context.get(Var { idx: 16 }), M31::from(11u32).into());
+    assert_eq!(context.get(Var { idx: 3 }), qm31_from_u32s(11, 11, 11, 11));
+    context.circuit.check_yields();
+    context.validate_circuit();
+}