Usability improvements: Arbitrary sized vectors support (#213)

* encrypted vector storage update

* bump version

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.3.0a4
+current_version = 0.3.0a5
 commit = True
 tag = True
 files = tenseal/version.py

tenseal/binding.cpp

@@ -1,3 +1,4 @@
+#include <pybind11/iostream.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -80,12 +81,16 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
                                                       py::module_local())
         .def(py::init([](const shared_ptr<TenSEALContext> &ctx,
                          const vector<int64_t> &data) {
-            return BFVVector::Create(ctx, data);
-        }))
-        .def(py::init(
-            [](const shared_ptr<TenSEALContext> &ctx, const std::string &data) {
-                return BFVVector::Create(ctx, data);
-            }))
+                 return BFVVector::Create(ctx, data);
+             }),
+             py::call_guard<py::scoped_ostream_redirect,
+                            py::scoped_estream_redirect>())
+        .def(py::init([](const shared_ptr<TenSEALContext> &ctx,
+                         const std::string &data) {
+                 return BFVVector::Create(ctx, data);
+             }),
+             py::call_guard<py::scoped_ostream_redirect,
+                            py::scoped_estream_redirect>())
         .def(py::init(
             [](const std::string &data) { return BFVVector::Create(data); }))
         .def("size", py::overload_cast<>(&BFVVector::size, py::const_))
@@ -215,17 +220,23 @@ PYBIND11_MODULE(_tenseal_cpp, m) {
         // specifying scale
         .def(py::init([](const shared_ptr<TenSEALContext> &ctx,
                          const vector<double> &data, double scale) {
-            return CKKSVector::Create(ctx, data, scale);
-        }))
+                 return CKKSVector::Create(ctx, data, scale);
+             }),
+             py::call_guard<py::scoped_ostream_redirect,
+                            py::scoped_estream_redirect>())
         // using global_scale if set
         .def(py::init([](const shared_ptr<TenSEALContext> &ctx,
                          const vector<double> &data) {
-            return CKKSVector::Create(ctx, data);
-        }))
-        .def(py::init(
-            [](const shared_ptr<TenSEALContext> &ctx, const std::string &data) {
-                return CKKSVector::Create(ctx, data);
-            }))
+                 return CKKSVector::Create(ctx, data);
+             }),
+             py::call_guard<py::scoped_ostream_redirect,
+                            py::scoped_estream_redirect>())
+        .def(py::init([](const shared_ptr<TenSEALContext> &ctx,
+                         const std::string &data) {
+                 return CKKSVector::Create(ctx, data);
+             }),
+             py::call_guard<py::scoped_ostream_redirect,
+                            py::scoped_estream_redirect>())
         .def(py::init(
             [](const std::string &data) { return CKKSVector::Create(data); }))
         .def("size", py::overload_cast<>(&CKKSVector::size, py::const_))

tenseal/cpp/tensors/bfvvector.cpp

@@ -14,9 +14,30 @@ void BFVVector::prepare_context(const shared_ptr<TenSEALContext>& ctx) {
 BFVVector::BFVVector(const shared_ptr<TenSEALContext>& ctx,
                      const plain_t& vec) {
     this->prepare_context(ctx);
-    // Encrypts the whole vector into a single ciphertext using BFV batching
-    this->_ciphertext = BFVVector::encrypt(ctx, vec);
-    this->_size = vec.size();
+
+    if (vec.empty()) {
+        throw invalid_argument("Attempting to encrypt an empty vector");
+    }
+
+    auto slot_count = ctx->slot_count<BatchEncoder>();
+    auto vec_chunks = vec.chunks(slot_count);
+
+    if (vec_chunks.size() > 1) {
+        std::cout
+            << "WARNING: The input does not fit in a single ciphertext, and "
+               "some operations will be disabled.\n"
+               "The following operations are disabled in this setup: matmul, "
+               "matmul_plain, conv2d_im2col, replicate_first_slot.\n"
+               "If you need to use those operations, try increasing the "
+               "poly_modulus parameter, to fit your input.\n";
+    }
+    this->_ciphertexts = vector<Ciphertext>();
+    this->_sizes = vector<size_t>();
+
+    for (auto& chunk : vec_chunks) {
+        this->_ciphertexts.push_back(BFVVector::encrypt(ctx, chunk));
+        this->_sizes.push_back(chunk.size());
+    }
 }
 
 BFVVector::BFVVector(const shared_ptr<TenSEALContext>& ctx, const string& vec) {
@@ -39,8 +60,8 @@ BFVVector::BFVVector(const TenSEALContextProto& ctx,
 }
 BFVVector::BFVVector(const shared_ptr<const BFVVector>& vec) {
     this->prepare_context(vec->tenseal_context());
-    this->_size = vec->size();
-    this->_ciphertext = vec->ciphertext();
+    this->_sizes = vec->chunked_size();
+    this->_ciphertexts = vec->ciphertext();
 }
 
 Ciphertext BFVVector::encrypt(shared_ptr<TenSEALContext> context,
@@ -65,15 +86,27 @@ Ciphertext BFVVector::encrypt(shared_ptr<TenSEALContext> context,
 }
 
 BFVVector::plain_t BFVVector::decrypt(const shared_ptr<SecretKey>& sk) const {
-    Plaintext plaintext;
     vector<int64_t> result;
+    result.reserve(this->size());
+
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        vector<int64_t> partial_result;
+
+        Plaintext plaintext;
+        this->tenseal_context()->decrypt(*sk, this->_ciphertexts[idx],
+                                         plaintext);
+        this->tenseal_context()->decode<BatchEncoder>(plaintext,
+                                                      partial_result);
+
+        // result contains all slots of ciphertext (poly_modulus_degree)
+        // we use the real vector size to delimit the resulting plaintext vector
+        auto partial_decr =
+            vector<int64_t>(partial_result.cbegin(),
+                            partial_result.cbegin() + this->_sizes[idx]);
+        result.insert(result.end(), partial_decr.begin(), partial_decr.end());
+    }
 
-    this->tenseal_context()->decrypt(*sk, this->_ciphertext, plaintext);
-    this->tenseal_context()->decode<BatchEncoder>(plaintext, result);
-
-    // result contains all slots of ciphertext (poly_modulus_degree)
-    // we use the real vector size to delimit the resulting plaintext vector
-    return vector<int64_t>(result.cbegin(), result.cbegin() + this->size());
+    return result;
 }
 
 shared_ptr<BFVVector> BFVVector::power_inplace(unsigned int power) {
@@ -106,14 +139,17 @@ shared_ptr<BFVVector> BFVVector::power_inplace(unsigned int power) {
 }
 
 shared_ptr<BFVVector> BFVVector::negate_inplace() {
-    this->tenseal_context()->evaluator->negate_inplace(this->_ciphertext);
+    for (auto& ct : this->_ciphertexts)
+        this->tenseal_context()->evaluator->negate_inplace(ct);
 
     return shared_from_this();
 }
 
 shared_ptr<BFVVector> BFVVector::square_inplace() {
-    this->tenseal_context()->evaluator->square_inplace(this->_ciphertext);
-    this->auto_relin(_ciphertext);
+    for (auto& ct : this->_ciphertexts) {
+        this->tenseal_context()->evaluator->square_inplace(ct);
+        this->auto_relin(ct);
+    }
 
     return shared_from_this();
 }
@@ -129,8 +165,9 @@ shared_ptr<BFVVector> BFVVector::add_inplace(
 
     this->broadcast_or_throw(to_add);
 
-    this->tenseal_context()->evaluator->add_inplace(this->_ciphertext,
-                                                    to_add->_ciphertext);
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx)
+        this->tenseal_context()->evaluator->add_inplace(
+            this->_ciphertexts[idx], to_add->_ciphertexts[idx]);
 
     return shared_from_this();
 }
@@ -146,8 +183,9 @@ shared_ptr<BFVVector> BFVVector::sub_inplace(
 
     this->broadcast_or_throw(to_sub);
 
-    this->tenseal_context()->evaluator->sub_inplace(this->_ciphertext,
-                                                    to_sub->_ciphertext);
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx)
+        this->tenseal_context()->evaluator->sub_inplace(
+            this->_ciphertexts[idx], to_sub->_ciphertexts[idx]);
 
     return shared_from_this();
 }
@@ -163,9 +201,11 @@ shared_ptr<BFVVector> BFVVector::mul_inplace(
 
     this->broadcast_or_throw(to_mul);
 
-    this->tenseal_context()->evaluator->multiply_inplace(this->_ciphertext,
-                                                         to_mul->_ciphertext);
-    this->auto_relin(_ciphertext);
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        this->tenseal_context()->evaluator->multiply_inplace(
+            this->_ciphertexts[idx], to_mul->_ciphertexts[idx]);
+        this->auto_relin(_ciphertexts[idx]);
+    }
 
     return shared_from_this();
 }
@@ -187,8 +227,18 @@ shared_ptr<BFVVector> BFVVector::dot_plain_inplace(
 }
 
 shared_ptr<BFVVector> BFVVector::sum_inplace(size_t /*axis=0*/) {
-    sum_vector(this->tenseal_context(), this->_ciphertext, this->size());
-    this->_size = 1;
+    vector<Ciphertext> interm_sum;
+    // TODO use multithreading for sum
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        Ciphertext out = this->_ciphertexts[idx];
+        sum_vector(this->tenseal_context(), out, this->_sizes[idx]);
+        interm_sum.push_back(out);
+    }
+    Ciphertext result;
+    tenseal_context()->evaluator->add_many(interm_sum, result);
+
+    this->_ciphertexts = {result};
+    this->_sizes = {1};
     return shared_from_this();
 }
 
@@ -198,17 +248,20 @@ shared_ptr<BFVVector> BFVVector::add_plain_inplace(
 }
 
 shared_ptr<BFVVector> BFVVector::add_plain_inplace(
-    const BFVVector::plain_t& to_add) {
-    if (this->size() != to_add.size()) {
+    const BFVVector::plain_t& vector_to_add) {
+    if (this->size() != vector_to_add.size()) {
         throw invalid_argument("can't add vectors of different sizes");
     }
-
-    Plaintext plaintext;
-
-    this->tenseal_context()->encode<BatchEncoder>(to_add.data_ref(), plaintext);
-    this->tenseal_context()->evaluator->add_plain_inplace(this->_ciphertext,
-                                                          plaintext);
-
+    auto slot_count = tenseal_context()->slot_count<BatchEncoder>();
+    auto to_add = vector_to_add.chunks(slot_count);
+
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        Plaintext plaintext;
+        this->tenseal_context()->encode<BatchEncoder>(to_add[idx].data_ref(),
+                                                      plaintext);
+        this->tenseal_context()->evaluator->add_plain_inplace(
+            this->_ciphertexts[idx], plaintext);
+    }
     return shared_from_this();
 }
 
@@ -218,17 +271,22 @@ shared_ptr<BFVVector> BFVVector::sub_plain_inplace(
 }
 
 shared_ptr<BFVVector> BFVVector::sub_plain_inplace(
-    const BFVVector::plain_t& to_sub) {
-    if (this->size() != to_sub.size()) {
+    const BFVVector::plain_t& vector_to_sub) {
+    if (this->size() != vector_to_sub.size()) {
         throw invalid_argument("can't sub vectors of different sizes");
     }
 
-    Plaintext plaintext;
+    auto slot_count = tenseal_context()->slot_count<BatchEncoder>();
+    auto to_sub = vector_to_sub.chunks(slot_count);
 
-    this->tenseal_context()->encode<BatchEncoder>(to_sub.data_ref(), plaintext);
-    this->tenseal_context()->evaluator->sub_plain_inplace(this->_ciphertext,
-                                                          plaintext);
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        Plaintext plaintext;
 
+        this->tenseal_context()->encode<BatchEncoder>(to_sub[idx].data_ref(),
+                                                      plaintext);
+        this->tenseal_context()->evaluator->sub_plain_inplace(
+            this->_ciphertexts[idx], plaintext);
+    }
     return shared_from_this();
 }
 
@@ -238,23 +296,29 @@ shared_ptr<BFVVector> BFVVector::mul_plain_inplace(
 }
 
 shared_ptr<BFVVector> BFVVector::mul_plain_inplace(
-    const BFVVector::plain_t& to_mul) {
-    if (this->size() != to_mul.size()) {
+    const BFVVector::plain_t& vector_to_mul) {
+    if (this->size() != vector_to_mul.size()) {
         throw invalid_argument("can't multiply vectors of different sizes");
     }
 
-    Plaintext plaintext;
-    this->tenseal_context()->encode<BatchEncoder>(to_mul.data_ref(), plaintext);
-
-    try {
-        this->tenseal_context()->evaluator->multiply_plain_inplace(
-            this->_ciphertext, plaintext);
-    } catch (const std::logic_error& e) {
-        if (strcmp(e.what(), "result ciphertext is transparent") == 0) {
-            // replace by encryption of zero
-            this->tenseal_context()->encrypt_zero(this->_ciphertext);
-        } else {  // Something else, need to be forwarded
-            throw;
+    auto slot_count = tenseal_context()->slot_count<BatchEncoder>();
+    auto to_mul = vector_to_mul.chunks(slot_count);
+
+    for (size_t idx = 0; idx < this->_ciphertexts.size(); ++idx) {
+        Plaintext plaintext;
+        this->tenseal_context()->encode<BatchEncoder>(to_mul[idx].data_ref(),
+                                                      plaintext);
+
+        try {
+            this->tenseal_context()->evaluator->multiply_plain_inplace(
+                this->_ciphertexts[idx], plaintext);
+        } catch (const std::logic_error& e) {
+            if (strcmp(e.what(), "result ciphertext is transparent") == 0) {
+                // replace by encryption of zero
+                this->tenseal_context()->encrypt_zero(this->_ciphertexts[idx]);
+            } else {  // Something else, need to be forwarded
+                throw;
+            }
         }
     }
 
@@ -282,40 +346,52 @@ shared_ptr<BFVVector> BFVVector::enc_matmul_plain_inplace(
 }
 
 shared_ptr<BFVVector> BFVVector::replicate_first_slot_inplace(size_t n) {
+    if (this->_ciphertexts.size() != 1)
+        throw invalid_argument(
+            "can't execute replicate_first_slot on chunked vectors");
+
     // mask
-    vector<int64_t> mask(this->_size, 0);
+    vector<int64_t> mask(this->size(), 0);
     mask[0] = 1;
     this->mul_plain_inplace(mask);
 
     // replicate
-    Ciphertext tmp = this->_ciphertext;
+    Ciphertext tmp = this->_ciphertexts[0];
     auto galois_keys = this->tenseal_context()->galois_keys();
     for (size_t i = 0; i < (size_t)ceil(log2(n)); i++) {
         this->tenseal_context()->evaluator->rotate_vector_inplace(
             tmp, static_cast<int>(-pow(2, i)), *galois_keys);
-        this->tenseal_context()->evaluator->add_inplace(this->_ciphertext, tmp);
-        tmp = this->_ciphertext;
+        this->tenseal_context()->evaluator->add_inplace(this->_ciphertexts[0],
+                                                        tmp);
+        tmp = this->_ciphertexts[0];
     }
 
-    this->_size = n;
+    this->_sizes = {n};
     return shared_from_this();
 }
 
 void BFVVector::load_proto(const BFVVectorProto& vec) {
     if (this->tenseal_context() == nullptr) {
         throw invalid_argument("context missing for deserialization");
     }
-    this->_size = vec.size();
-    this->_ciphertext = SEALDeserialize<Ciphertext>(
-        *this->tenseal_context()->seal_context(), vec.ciphertext());
+    this->_sizes = vector<size_t>();
+    this->_ciphertexts = vector<Ciphertext>();
+
+    for (auto& sz : vec.sizes()) this->_sizes.push_back(sz);
+    for (auto& ct : vec.ciphertexts())
+        this->_ciphertexts.push_back(SEALDeserialize<Ciphertext>(
+            *this->tenseal_context()->seal_context(), ct));
 }
 
 BFVVectorProto BFVVector::save_proto() const {
     BFVVectorProto buffer;
 
-    *buffer.mutable_ciphertext() = SEALSerialize<Ciphertext>(this->_ciphertext);
-    buffer.set_size(static_cast<int>(this->_size));
-
+    for (auto& ct : this->_ciphertexts) {
+        buffer.add_ciphertexts(SEALSerialize<Ciphertext>(ct));
+    }
+    for (auto& sz : this->_sizes) {
+        buffer.add_sizes(sz);
+    }
     return buffer;
 }