stwo 1.0.0

Stwo Circle-STARKs prover
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use core::ops::Sub;
use std::iter::zip;
use std::ops::{Add, Mul};

use crate::core::fields::m31::BaseField;
use crate::core::fields::qm31::SecureField;
use crate::prover::backend::simd::column::SecureColumn;
use crate::prover::backend::simd::m31::N_LANES;
use crate::prover::backend::simd::qm31::PackedSecureField;
use crate::prover::backend::simd::SimdBackend;
use crate::prover::backend::{Column, CpuBackend};
use crate::prover::lookups::mle::{Mle, MleOps};

impl MleOps<BaseField> for SimdBackend {
    fn fix_first_variable(
        mle: Mle<Self, BaseField>,
        assignment: SecureField,
    ) -> Mle<Self, SecureField> {
        let midpoint = mle.len() / 2;

        // Use CPU backend to avoid dealing with instances smaller than `PackedSecureField`.
        if midpoint < N_LANES {
            let cpu_mle = Mle::<CpuBackend, BaseField>::new(mle.to_cpu());
            let cpu_res = cpu_mle.fix_first_variable(assignment);
            return Mle::new(cpu_res.into_evals().into_iter().collect());
        }

        let packed_assignment = PackedSecureField::broadcast(assignment);
        let packed_midpoint = midpoint / N_LANES;
        let (evals_at_0x, evals_at_1x) = mle.data.split_at(packed_midpoint);

        let res = zip(evals_at_0x, evals_at_1x)
            // MLE at points `({0, 1}, rev(bits(i)), v)` for all `v` in `{0, 1}^LOG_N_SIMD_LANES`.
            .map(|(&packed_eval_at_0iv, &packed_eval_at_1iv)| {
                fold_packed_mle_evals(packed_assignment, packed_eval_at_0iv, packed_eval_at_1iv)
            })
            .collect();

        Mle::new(res)
    }
}

impl MleOps<SecureField> for SimdBackend {
    fn fix_first_variable(
        mle: Mle<Self, SecureField>,
        assignment: SecureField,
    ) -> Mle<Self, SecureField> {
        let midpoint = mle.len() / 2;

        // Use CPU backend to avoid dealing with instances smaller than `PackedSecureField`.
        if midpoint < N_LANES {
            let cpu_mle = Mle::<CpuBackend, SecureField>::new(mle.to_cpu());
            let cpu_res = cpu_mle.fix_first_variable(assignment);
            return Mle::new(cpu_res.into_evals().into_iter().collect());
        }

        let packed_midpoint = midpoint / N_LANES;
        let packed_assignment = PackedSecureField::broadcast(assignment);
        let mut packed_evals = mle.into_evals().data;

        for i in 0..packed_midpoint {
            // MLE at points `({0, 1}, rev(bits(i)), v)` for all `v` in `{0, 1}^LOG_N_SIMD_LANES`.
            let packed_eval_at_0iv = packed_evals[i];
            let packed_eval_at_1iv = packed_evals[i + packed_midpoint];
            packed_evals[i] =
                fold_packed_mle_evals(packed_assignment, packed_eval_at_0iv, packed_eval_at_1iv);
        }

        packed_evals.truncate(packed_midpoint);

        let length = packed_evals.len() * N_LANES;
        let data = packed_evals;

        Mle::new(SecureColumn { data, length })
    }
}

/// Computes all `eq(0, assignment_i) * eval0_i + eq(1, assignment_i) * eval1_i`.
// TODO(andrew): Remove complex trait bounds once we have something like
// AbstractField/AbstractExtensionField traits.
fn fold_packed_mle_evals<
    PackedF: Sub<Output = PackedF> + Copy,
    PackedEF: Mul<PackedF, Output = PackedEF> + Add<PackedF, Output = PackedEF>,
>(
    assignment: PackedEF,
    eval0: PackedF,
    eval1: PackedF,
) -> PackedEF {
    assignment * (eval1 - eval0) + eval0
}

#[cfg(test)]
mod tests {
    use itertools::Itertools;

    use crate::core::channel::Channel;
    use crate::core::fields::m31::BaseField;
    use crate::core::fields::qm31::SecureField;
    use crate::core::test_utils::test_channel;
    use crate::prover::backend::simd::SimdBackend;
    use crate::prover::backend::{Column, CpuBackend};
    use crate::prover::lookups::mle::Mle;

    #[test]
    fn fix_first_variable_with_secure_field_mle_matches_cpu() {
        const N_VARIABLES: u32 = 8;
        let values = test_channel().draw_secure_felts(1 << N_VARIABLES);
        let mle_simd = Mle::<SimdBackend, SecureField>::new(values.iter().copied().collect());
        let mle_cpu = Mle::<CpuBackend, SecureField>::new(values);
        let random_assignment = SecureField::from_u32_unchecked(7, 12, 3, 2);
        let mle_fixed_cpu = mle_cpu.fix_first_variable(random_assignment);

        let mle_fixed_simd = mle_simd.fix_first_variable(random_assignment);

        assert_eq!(mle_fixed_simd.into_evals().to_cpu(), *mle_fixed_cpu)
    }

    #[test]
    fn fix_first_variable_with_base_field_mle_matches_cpu() {
        const N_VARIABLES: u32 = 8;
        let values = (0..1 << N_VARIABLES).map(BaseField::from).collect_vec();
        let mle_simd = Mle::<SimdBackend, BaseField>::new(values.iter().copied().collect());
        let mle_cpu = Mle::<CpuBackend, BaseField>::new(values);
        let random_assignment = SecureField::from_u32_unchecked(7, 12, 3, 2);
        let mle_fixed_cpu = mle_cpu.fix_first_variable(random_assignment);

        let mle_fixed_simd = mle_simd.fix_first_variable(random_assignment);

        assert_eq!(mle_fixed_simd.into_evals().to_cpu(), *mle_fixed_cpu)
    }
}