Skip to main content

poulpy_core/api/
transfer.rs

1//! Moving a layout from one backend to another.
2//!
3//! The destination must already exist:
4//!
5//! ```ignore
6//! let mut sk = module.glwe_secret_alloc_from_infos(&sk_host);
7//! sk_host.transfer_into(&mut sk);
8//! ```
9//!
10//! Three things follow from requiring an allocated destination, none of which a
11//! form returning a fresh value can offer:
12//!
13//! - no backend is named anywhere. Both buffer types are concrete, so nothing
14//!   has to be recovered from an associated-type projection, and there is no
15//!   type parameter whose direction a reader can misjudge.
16//! - the whole shape is checkable, not just the byte count. Both operands are
17//!   present, so `base2k`, `k`, degree, column and limb counts are all compared.
18//! - the allocation is visible and hoistable. Moving a key set allocates once,
19//!   and the cost is not hidden inside something that reads like a conversion.
20//!
21//! Impls are one per layout: a shape check, then the container move. Being a
22//! trait rather than a closed set of methods, a downstream crate can transfer
23//! its own layouts without editing this one.
24
25use poulpy_hal::layouts::{
26    CopyFromHost, CopyToHost, Data, DataView, DataViewMut, MatZnx, ScalarZnx, VecZnx, ZnxWord, transfer_buf_into,
27};
28
29use crate::layouts::{
30    GGLWE, GGLWEToGGSWKey, GGSW, GLWE, GLWEAutomorphismKey, GLWEPlaintext, GLWESecret, GLWESwitchingKey, GLWETensor,
31    GLWETensorKey, GLWEToLWEKey, LWE, LWEPlaintext, LWESecret,
32};
33
34/// Moves `Self` into an already-allocated destination.
35pub trait TransferInto<Dst> {
36    /// # Panics
37    ///
38    /// If the two values do not agree on shape.
39    fn transfer_into(&self, dst: &mut Dst);
40}
41
42fn move_vec_znx<D1, D2, W>(src: &VecZnx<D1, W>, dst: &mut VecZnx<D2, W>)
43where
44    D1: Data + CopyToHost,
45    D2: Data + CopyFromHost,
46    W: ZnxWord,
47{
48    assert_eq!(src.n(), dst.n(), "transfer_into: ring degree");
49    assert_eq!(src.cols(), dst.cols(), "transfer_into: cols");
50    assert_eq!(src.size(), dst.size(), "transfer_into: size");
51    transfer_buf_into(src.data(), dst.data_mut());
52}
53
54fn move_mat_znx<D1, D2, W>(src: &MatZnx<D1, W>, dst: &mut MatZnx<D2, W>)
55where
56    D1: Data + CopyToHost,
57    D2: Data + CopyFromHost,
58    W: ZnxWord,
59{
60    assert_eq!(src.n(), dst.n(), "transfer_into: ring degree");
61    assert_eq!(src.rows(), dst.rows(), "transfer_into: rows");
62    assert_eq!(src.cols_in(), dst.cols_in(), "transfer_into: cols_in");
63    assert_eq!(src.cols_out(), dst.cols_out(), "transfer_into: cols_out");
64    assert_eq!(src.size(), dst.size(), "transfer_into: size");
65    transfer_buf_into(src.data(), dst.data_mut());
66}
67
68fn move_scalar_znx<D1, D2, W>(src: &ScalarZnx<D1, W>, dst: &mut ScalarZnx<D2, W>)
69where
70    D1: Data + CopyToHost,
71    D2: Data + CopyFromHost,
72    W: ZnxWord,
73{
74    assert_eq!(src.n(), dst.n(), "transfer_into: ring degree");
75    assert_eq!(src.cols(), dst.cols(), "transfer_into: cols");
76    transfer_buf_into(src.data(), dst.data_mut());
77}
78
79impl<D1, D2, W> TransferInto<GLWE<D2, W>> for GLWE<D1, W>
80where
81    D1: Data + CopyToHost,
82    D2: Data + CopyFromHost,
83    W: ZnxWord,
84{
85    fn transfer_into(&self, dst: &mut GLWE<D2, W>) {
86        assert_eq!(self.base2k, dst.base2k, "transfer_into: GLWE base2k");
87        assert_eq!(self.k, dst.k, "transfer_into: GLWE k");
88        move_vec_znx(&self.data, &mut dst.data);
89    }
90}
91
92impl<D1, D2, W> TransferInto<GLWEPlaintext<D2, W>> for GLWEPlaintext<D1, W>
93where
94    D1: Data + CopyToHost,
95    D2: Data + CopyFromHost,
96    W: ZnxWord,
97{
98    fn transfer_into(&self, dst: &mut GLWEPlaintext<D2, W>) {
99        assert_eq!(self.base2k, dst.base2k, "transfer_into: GLWEPlaintext base2k");
100        assert_eq!(self.k, dst.k, "transfer_into: GLWEPlaintext k");
101        move_vec_znx(&self.data, &mut dst.data);
102    }
103}
104
105impl<D1, D2, W> TransferInto<LWEPlaintext<D2, W>> for LWEPlaintext<D1, W>
106where
107    D1: Data + CopyToHost,
108    D2: Data + CopyFromHost,
109    W: ZnxWord,
110{
111    fn transfer_into(&self, dst: &mut LWEPlaintext<D2, W>) {
112        assert_eq!(self.base2k, dst.base2k, "transfer_into: LWEPlaintext base2k");
113        assert_eq!(self.k, dst.k, "transfer_into: LWEPlaintext k");
114        move_vec_znx(&self.data, &mut dst.data);
115    }
116}
117
118impl<D1, D2, W> TransferInto<LWE<D2, W>> for LWE<D1, W>
119where
120    D1: Data + CopyToHost,
121    D2: Data + CopyFromHost,
122    W: ZnxWord,
123{
124    fn transfer_into(&self, dst: &mut LWE<D2, W>) {
125        assert_eq!(self.base2k, dst.base2k, "transfer_into: LWE base2k");
126        assert_eq!(self.k, dst.k, "transfer_into: LWE k");
127        move_vec_znx(&self.body, &mut dst.body);
128        move_vec_znx(&self.mask, &mut dst.mask);
129    }
130}
131
132impl<D1, D2, W> TransferInto<GGLWE<D2, W>> for GGLWE<D1, W>
133where
134    D1: Data + CopyToHost,
135    D2: Data + CopyFromHost,
136    W: ZnxWord,
137{
138    fn transfer_into(&self, dst: &mut GGLWE<D2, W>) {
139        assert_eq!(self.base2k, dst.base2k, "transfer_into: GGLWE base2k");
140        assert_eq!(self.k_aux, dst.k_aux, "transfer_into: GGLWE k_aux");
141        assert_eq!(self.dsize, dst.dsize, "transfer_into: GGLWE dsize");
142        move_mat_znx(&self.data, &mut dst.data);
143    }
144}
145
146impl<D1, D2, W> TransferInto<GGSW<D2, W>> for GGSW<D1, W>
147where
148    D1: Data + CopyToHost,
149    D2: Data + CopyFromHost,
150    W: ZnxWord,
151{
152    fn transfer_into(&self, dst: &mut GGSW<D2, W>) {
153        assert_eq!(self.base2k, dst.base2k, "transfer_into: GGSW base2k");
154        assert_eq!(self.k_aux, dst.k_aux, "transfer_into: GGSW k_aux");
155        assert_eq!(self.dsize, dst.dsize, "transfer_into: GGSW dsize");
156        move_mat_znx(&self.data, &mut dst.data);
157    }
158}
159
160impl<D1, D2, W> TransferInto<GLWESwitchingKey<D2, W>> for GLWESwitchingKey<D1, W>
161where
162    D1: Data + CopyToHost,
163    D2: Data + CopyFromHost,
164    W: ZnxWord,
165{
166    fn transfer_into(&self, dst: &mut GLWESwitchingKey<D2, W>) {
167        self.key.transfer_into(&mut dst.key);
168        dst.input_degree = self.input_degree;
169        dst.output_degree = self.output_degree;
170    }
171}
172
173impl<D1, D2, W> TransferInto<GLWEAutomorphismKey<D2, W>> for GLWEAutomorphismKey<D1, W>
174where
175    D1: Data + CopyToHost,
176    D2: Data + CopyFromHost,
177    W: ZnxWord,
178{
179    fn transfer_into(&self, dst: &mut GLWEAutomorphismKey<D2, W>) {
180        self.key.transfer_into(&mut dst.key);
181        dst.p = self.p;
182    }
183}
184
185impl<D1, D2, W> TransferInto<GLWESecret<D2, W>> for GLWESecret<D1, W>
186where
187    D1: Data + CopyToHost,
188    D2: Data + CopyFromHost,
189    W: ZnxWord,
190{
191    fn transfer_into(&self, dst: &mut GLWESecret<D2, W>) {
192        move_scalar_znx(&self.data, &mut dst.data);
193        dst.dist = self.dist;
194    }
195}
196
197impl<D1, D2, W> TransferInto<LWESecret<D2, W>> for LWESecret<D1, W>
198where
199    D1: Data + CopyToHost,
200    D2: Data + CopyFromHost,
201    W: ZnxWord,
202{
203    fn transfer_into(&self, dst: &mut LWESecret<D2, W>) {
204        move_scalar_znx(&self.data, &mut dst.data);
205        dst.dist = self.dist;
206    }
207}
208
209impl<D1, D2, W> TransferInto<GLWETensor<D2, W>> for GLWETensor<D1, W>
210where
211    D1: Data + CopyToHost,
212    D2: Data + CopyFromHost,
213    W: ZnxWord,
214{
215    fn transfer_into(&self, dst: &mut GLWETensor<D2, W>) {
216        assert_eq!(self.base2k, dst.base2k, "transfer_into: GLWETensor base2k");
217        assert_eq!(self.k, dst.k, "transfer_into: GLWETensor k");
218        assert_eq!(self.rank, dst.rank, "transfer_into: GLWETensor rank");
219        move_vec_znx(&self.data, &mut dst.data);
220    }
221}
222
223impl<D1, D2, W> TransferInto<GLWETensorKey<D2, W>> for GLWETensorKey<D1, W>
224where
225    D1: Data + CopyToHost,
226    D2: Data + CopyFromHost,
227    W: ZnxWord,
228{
229    fn transfer_into(&self, dst: &mut GLWETensorKey<D2, W>) {
230        self.0.transfer_into(&mut dst.0);
231    }
232}
233
234impl<D1, D2, W> TransferInto<GLWEToLWEKey<D2, W>> for GLWEToLWEKey<D1, W>
235where
236    D1: Data + CopyToHost,
237    D2: Data + CopyFromHost,
238    W: ZnxWord,
239{
240    fn transfer_into(&self, dst: &mut GLWEToLWEKey<D2, W>) {
241        self.0.transfer_into(&mut dst.0);
242    }
243}
244
245impl<D1, D2, W> TransferInto<GGLWEToGGSWKey<D2, W>> for GGLWEToGGSWKey<D1, W>
246where
247    D1: Data + CopyToHost,
248    D2: Data + CopyFromHost,
249    W: ZnxWord,
250{
251    fn transfer_into(&self, dst: &mut GGLWEToGGSWKey<D2, W>) {
252        assert_eq!(self.keys.len(), dst.keys.len(), "transfer_into: GGLWEToGGSWKey key count");
253        for (src, dst) in self.keys.iter().zip(&mut dst.keys) {
254            src.transfer_into(dst);
255        }
256    }
257}