poulpy_core/default/linear_transformation/
prepare.rs

1//! Transform preparation reference implementations.
2//!
3//! Implements docs/linear_transformation.md: diagonals are encoded by the scheme-aware
4//! caller (CKKS), then turned into right convolution operands (`CnvPVecR`).
5//! The resident transform is allocated up-front as a
6//! `LinearTransformation<PreparedDiagonal<…>>` via
7//! [`LinearTransformation::alloc_prepared`] from a [`LinearTransformationLayout`]
8//! and a plaintext-shape proxy; this module's `_into` function only fills the
9//! pre-allocated `CnvPVecR` slots, performing zero `CnvPVecR` allocations.
10//! Backends forward to them from their [`crate::oep::LinearTransformationDefault`]
11//! impl.
12
13use poulpy_hal::layouts::CnvPVecRToBackendMut;
14use poulpy_hal::{
15    api::{CnvPVecAlloc, Convolution},
16    layouts::{Backend, ScratchArena},
17};
18
19use crate::layouts::IntPolyInfos;
20use crate::{
21    default::operations::msb_mask_bottom_limb,
22    layouts::{
23        GLWEInfos, GLWEToBackendRef, LWEInfos, LinearTransformation, LinearTransformationDiagonal, LinearTransformationGiantStep,
24        LinearTransformationLayout, LinearTransformationPlan, prepared::PreparedDiagonal,
25    },
26};
27
28impl<BE: Backend> LinearTransformation<PreparedDiagonal<BE::OwnedBuf, BE>> {
29    /// Pre-allocates a resident (prepared) linear transformation sized for the
30    /// given BSGS `layout` and plaintext shape `pt_infos`.
31    ///
32    /// Convenience for the layout-driven flow: builds the BSGS index via
33    /// `layout.index()` and forwards to [`Self::alloc_prepared_from_index`].
34    pub fn alloc_prepared<M, P>(module: &M, layout: &LinearTransformationLayout, pt_infos: &P) -> Self
35    where
36        M: CnvPVecAlloc<BE>,
37        P: LWEInfos,
38    {
39        Self::alloc_prepared_from_index(module, &layout.index(), pt_infos)
40    }
41
42    /// Pre-allocates a resident transform sized for an explicit BSGS `index`.
43    ///
44    /// Each diagonal carries the plaintext's `base2k` / `k` so the evaluator
45    /// never needs the raw plaintext transform again; the convolution buffers are
46    /// zeroed and populated by `glwe_prepare_linear_transformation_rhs`. The
47    /// per-diagonal `log_scale` is left at `0` for the scheme layer to set.
48    pub fn alloc_prepared_from_index<M, P>(module: &M, index: &LinearTransformationPlan, pt_infos: &P) -> Self
49    where
50        M: CnvPVecAlloc<BE>,
51        P: LWEInfos,
52    {
53        let pt_size = pt_infos.size();
54        let base2k = pt_infos.base2k();
55        let k = pt_infos.k();
56
57        let mut giant_steps = Vec::with_capacity(index.giant_steps.len());
58        for (g, &rot) in index.giant_steps.iter().enumerate() {
59            let baby_rots = &index.index[g];
60            let mut diagonals = Vec::with_capacity(baby_rots.len());
61            for &baby in baby_rots {
62                diagonals.push(LinearTransformationDiagonal {
63                    baby,
64                    plaintext: PreparedDiagonal {
65                        cnv: module.cnv_pvec_right_alloc(1, pt_size),
66                        base2k,
67                        k,
68                        log_scale: 0,
69                    },
70                });
71            }
72            giant_steps.push(LinearTransformationGiantStep { rot, diagonals });
73        }
74
75        LinearTransformation {
76            baby_steps: index.baby_steps.clone(),
77            giant_steps,
78        }
79    }
80
81    /// Sets the per-diagonal `log_scale` of every diagonal; called by the scheme
82    /// layer during the populate step (mirrors the streamed plaintext's
83    /// `log_delta`).
84    pub fn set_log_scale(&mut self, log_scale: usize) {
85        for gs in &mut self.giant_steps {
86            for d in &mut gs.diagonals {
87                d.plaintext.set_log_scale(log_scale);
88            }
89        }
90    }
91
92    /// Base-2 log of the plaintext scaling factor shared by every diagonal.
93    pub fn log_scale(&self) -> usize {
94        self.first_diagonal_plaintext()
95            .expect("prepared linear transformation has no diagonals")
96            .log_scale()
97    }
98}
99
100/// Reference impl: scratch bytes for `glwe_prepare_linear_transformation_rhs`.
101pub fn glwe_prepare_linear_transformation_rhs_tmp_bytes_default<BE, M, P>(module: &M, pt_infos: &P) -> usize
102where
103    BE: Backend,
104    M: Convolution<BE>,
105    P: LWEInfos,
106{
107    module.cnv_prepare_right_tmp_bytes(pt_infos.size(), pt_infos.size())
108}
109
110/// Reference impl: encodes every diagonal of `lt` into the matching
111/// pre-allocated `CnvPVecR` slot in `prepared`.
112///
113/// `prepared` must have been sized via
114/// [`LinearTransformation::alloc_prepared`](LinearTransformation::alloc_prepared)
115/// for the same BSGS schedule (giant rotations and baby rotations) that `lt`
116/// follows.
117pub fn glwe_prepare_linear_transformation_rhs_default<BE, M, P>(
118    module: &M,
119    prepared: &mut LinearTransformation<PreparedDiagonal<BE::OwnedBuf, BE>>,
120    lt: &LinearTransformation<P>,
121    scratch: &mut ScratchArena<'_, BE>,
122) where
123    BE: Backend,
124    M: CnvPVecAlloc<BE> + Convolution<BE>,
125    P: GLWEToBackendRef<BE> + GLWEInfos,
126{
127    if !lt.baby_steps.is_empty() {
128        assert_eq!(
129            lt.baby_steps.first(),
130            Some(&0),
131            "baby_steps must start with the identity rotation (0)"
132        );
133    }
134
135    let first = prepared
136        .first_diagonal_plaintext()
137        .expect("prepared linear transformation has no diagonals");
138    let pt_base2k = first.base2k();
139    let pt_k = first.k();
140    let pt_base2k_usize = pt_base2k.as_usize();
141    let pt_k_usize = pt_k.as_usize();
142    // The diagonal is an integer poly encoded across its full physical width
143    // (`max_k`), so the bottom-limb mask must span `max_k`, not the (possibly
144    // smaller) effective `k`, otherwise the low limb's data is truncated.
145    let mask = msb_mask_bottom_limb(pt_base2k_usize, first.encoded_k().as_usize());
146
147    for gs in &lt.giant_steps {
148        if gs.diagonals.is_empty() {
149            continue;
150        }
151        let prepared_gs = prepared
152            .giant_steps
153            .iter_mut()
154            .find(|p| p.rot == gs.rot)
155            .unwrap_or_else(|| panic!("prepared cache has no giant step for rotation {}", gs.rot));
156
157        for d in &gs.diagonals {
158            let plaintext = &d.plaintext;
159            assert_eq!(
160                plaintext.base2k(),
161                pt_base2k,
162                "linear transformation diagonal base2k does not match prepared cache"
163            );
164            assert_eq!(
165                plaintext.k(),
166                pt_k,
167                "linear transformation diagonal k does not match prepared cache"
168            );
169            assert_eq!(
170                pt_k_usize.div_ceil(pt_base2k_usize),
171                plaintext.size(),
172                "linear transformation plaintext size does not match its effective precision"
173            );
174
175            let prepared_slot = prepared_gs
176                .diagonals
177                .iter_mut()
178                .find(|p| p.baby == d.baby)
179                .unwrap_or_else(|| panic!("prepared cache has no diagonal slot for baby {} at giant {}", d.baby, gs.rot));
180            let plaintext_backend = plaintext.to_backend_ref();
181            module.cnv_prepare_right(
182                &mut prepared_slot.plaintext.cnv_mut().to_backend_mut(),
183                &plaintext_backend.data,
184                mask,
185                scratch,
186            );
187        }
188    }
189}
poulpy_core/default/linear_transformation/prepare.rs

poulpy_core/default/linear_transformation/
prepare.rs
