Repository: secretflow/spu (language: python)

Issue:
[Bug]: The critical condition judgment is wrong
### Issue Type

Usability

### Modules Involved

MPC protocol

### Have you reproduced the bug with SPU HEAD?

No

### Have you searched existing issues?

Yes

### SPU Version

07a5aa5890375c531bdfdb257c6e2a19dec800d7

### OS Platform and Distribution

Ubuntu 20.04

### Python Version

3.10

### Compiler Version

_No response_

### Current Behavior?

In libspu/mpc/cheetah/rlwe/modswitch_helper.cc, there is a bug in the function `template <typename Scalar>
  void CenteralizeAt(NdArrayView<const Scalar> src, size_t mod_idx,
                     absl::Span<uint64_t> out) const`.
The main problem is the for cycle
```
for (size_t i = 0; i < n; ++i) {
      auto x128 = static_cast<uint128_t>(src[i]);
      if (x128 > t_half_) {
        uint64_t u = BarrettReduce(-x128 & mod_t_mask_, mod_qj);
        out[i] = negate_uint_mod(u, mod_qj);
      } else {
        out[i] = BarrettReduce(src[i], mod_qj);
      }
    }
```
When `x128 = t_half_`, it shoudl go to the first branch instead of the second branch due to the annotation `// view x \in [0, 2^k) as [-2^{k-1}, 2^{k-1})`. Note that the two branches give different results for `x128 = t_half_`.

### Standalone code to reproduce the issue

```Python
template <typename Scalar>
  void CenteralizeAt(NdArrayView<const Scalar> src, size_t mod_idx,
                     absl::Span<uint64_t> out) const {
    using namespace seal::util;
    SPU_ENFORCE_EQ(sizeof(Scalar) * 8, absl::bit_ceil(base_mod_bitlen_),
                   fmt::format("expect base_mod_bitlen={} but got {}",
                               base_mod_bitlen_, sizeof(Scalar) * 8));

    const auto &modulus = context_.key_context_data()->parms().coeff_modulus();
    SPU_ENFORCE(mod_idx < coeff_modulus_size(), "Centeralize: invalid mod_idx");
    size_t n = src.numel();
    SPU_ENFORCE(n == out.size(), "Centeralize: size mismatch");

    const auto &mod_qj = modulus[mod_idx];
    // view x \in [0, 2^k) as [-2^{k-1}, 2^{k-1})
    for (size_t i = 0; i < n; ++i) {
      auto x128 = static_cast<uint128_t>(src[i]);
      if (x128 > t_half_) {
        uint64_t u = BarrettReduce(-x128 & mod_t_mask_, mod_qj);
        out[i] = negate_uint_mod(u, mod_qj);
      } else {
        out[i] = BarrettReduce(src[i], mod_qj);
      }
    }
  }
```


### Relevant log output

_No response_
[Bug]: The critical condition judgment is wrong
### Issue Type

Usability

### Modules Involved

MPC protocol

### Have you reproduced the bug with SPU HEAD?

No

### Have you searched existing issues?

Yes

### SPU Version

07a5aa5890375c531bdfdb257c6e2a19dec800d7

### OS Platform and Distribution

Ubuntu 20.04

### Python Version

3.10

### Compiler Version

_No response_

### Current Behavior?

In libspu/mpc/cheetah/rlwe/modswitch_helper.cc, there is a bug in the function `template <typename Scalar>
  void CenteralizeAt(NdArrayView<const Scalar> src, size_t mod_idx,
                     absl::Span<uint64_t> out) const`.
The main problem is the for cycle
```
for (size_t i = 0; i < n; ++i) {
      auto x128 = static_cast<uint128_t>(src[i]);
      if (x128 > t_half_) {
        uint64_t u = BarrettReduce(-x128 & mod_t_mask_, mod_qj);
        out[i] = negate_uint_mod(u, mod_qj);
      } else {
        out[i] = BarrettReduce(src[i], mod_qj);
      }
    }
```
When `x128 = t_half_`, it shoudl go to the first branch instead of the second branch due to the annotation `// view x \in [0, 2^k) as [-2^{k-1}, 2^{k-1})`. Note that the two branches give different results for `x128 = t_half_`.

### Standalone code to reproduce the issue

```Python
template <typename Scalar>
  void CenteralizeAt(NdArrayView<const Scalar> src, size_t mod_idx,
                     absl::Span<uint64_t> out) const {
    using namespace seal::util;
    SPU_ENFORCE_EQ(sizeof(Scalar) * 8, absl::bit_ceil(base_mod_bitlen_),
                   fmt::format("expect base_mod_bitlen={} but got {}",
                               base_mod_bitlen_, sizeof(Scalar) * 8));

    const auto &modulus = context_.key_context_data()->parms().coeff_modulus();
    SPU_ENFORCE(mod_idx < coeff_modulus_size(), "Centeralize: invalid mod_idx");
    size_t n = src.numel();
    SPU_ENFORCE(n == out.size(), "Centeralize: size mismatch");

    const auto &mod_qj = modulus[mod_idx];
    // view x \in [0, 2^k) as [-2^{k-1}, 2^{k-1})
    for (size_t i = 0; i < n; ++i) {
      auto x128 = static_cast<uint128_t>(src[i]);
      if (x128 > t_half_) {
        uint64_t u = BarrettReduce(-x128 & mod_t_mask_, mod_qj);
        out[i] = negate_uint_mod(u, mod_qj);
      } else {
        out[i] = BarrettReduce(src[i], mod_qj);
      }
    }
  }
```


### Relevant log output

_No response_


Discussion / hints:
@fionser , is this a bug?
Similar problems appear at `template <typename Scalar>
  void ModulusDownRNS(absl::Span<const uint64_t> src,
                      absl::Span<Scalar> out) const `:
```
          // [0, gamma) -> [-gamma/2, gamma/2]
          if (on_gamma > gamma_div_2) {
            return ((on_t + gamma_.value() - on_gamma) * inv_gamma_mod_t_) &
                   mod_t_mask_;
          } else {
            return ((on_t - on_gamma) * inv_gamma_mod_t_) & mod_t_mask_;
          }
```

More questions: 
(2) in packlwes.cc. The function `void NegacyclicRightShiftInplace(RLWECt &ct, size_t shift,
                                 const seal::SEALContext &context)` is just multiplying a plaintext $X^{shift}$. So it can be replaced by `seal::Evaluator:: multiply_plain`, right?
(3) at line 356 in packlwes_test.cc, the index in `xtmp[num_coeffs - 1 - i] = -xvec[i];` is not right. And the minus sign will cause wrong answer, since we want `-xvec[i] mod q[mod_idx]`, while `-xvec[i]` is computed module $2^{k}$.
- Viewing $x \in [0, 2^k)$ as $[-2^{k-1}, 2^{k-1})$ or in $(-2^{k-1}, 2^{k-1}]$ should be all right for `CenteralizeAt.
- Yes, the document is mismatching the code in `CenteralizeAt` where the code is actually doing $(-2^{k-1}, 2^{k-1}]$ 
- For the `[0, gamma) -> [-gamma/2, gamma/2]` since `gamma` is an odd value. 
- `NegacyclicRightShiftInplace` might be efficient than using `multiply_plain`. At least, you don't really need the plain polynomial except the non-zero index.
- `-xvec[i] mod 2^k` is fine, since that is what we what. The `mod q` part is defered to Line 365 or Line 367
 
(3) `xtmp[num_coeffs - 1 - i]` should be `xtmp[num_coeffs - i]` . Note that `-xvec[i]` mod $2^k$ mod q is different from `-xvec[i] mod q`. Actually, a possible case is `-xvec[i]` mod $2^k$ = $2^k$ - xvec[i]. And I do not understand the part $2^k$ mod q.
@fionser , is this a bug?
Similar problems appear at `template <typename Scalar>
  void ModulusDownRNS(absl::Span<const uint64_t> src,
                      absl::Span<Scalar> out) const `:
```
          // [0, gamma) -> [-gamma/2, gamma/2]
          if (on_gamma > gamma_div_2) {
            return ((on_t + gamma_.value() - on_gamma) * inv_gamma_mod_t_) &
                   mod_t_mask_;
          } else {
            return ((on_t - on_gamma) * inv_gamma_mod_t_) & mod_t_mask_;
          }
```

More questions: 
(2) in packlwes.cc. The function `void NegacyclicRightShiftInplace(RLWECt &ct, size_t shift,
                                 const seal::SEALContext &context)` is just multiplying a plaintext $X^{shift}$. So it can be replaced by `seal::Evaluator:: multiply_plain`, right?
(3) at line 356 in packlwes_test.cc, the index in `xtmp[num_coeffs - 1 - i] = -xvec[i];` is not right. And the minus sign will cause wrong answer, since we want `-xvec[i] mod q[mod_idx]`, while `-xvec[i]` is computed module $2^{k}$.
- Viewing $x \in [0, 2^k)$ as $[-2^{k-1}, 2^{k-1})$ or in $(-2^{k-1}, 2^{k-1}]$ should be all right for `CenteralizeAt.
- Yes, the document is mismatching the code in `CenteralizeAt` where the code is actually doing $(-2^{k-1}, 2^{k-1}]$ 
- For the `[0, gamma) -> [-gamma/2, gamma/2]` since `gamma` is an odd value. 
- `NegacyclicRightShiftInplace` might be efficient than using `multiply_plain`. At least, you don't really need the plain polynomial except the non-zero index.
- `-xvec[i] mod 2^k` is fine, since that is what we what. The `mod q` part is defered to Line 365 or Line 367
 
(3) `xtmp[num_coeffs - 1 - i]` should be `xtmp[num_coeffs - i]` . Note that `-xvec[i]` mod $2^k$ mod q is different from `-xvec[i] mod q`. Actually, a possible case is `-xvec[i]` mod $2^k$ = $2^k$ - xvec[i]. And I do not understand the part $2^k$ mod q.

Fix patch:
diff --git a/libspu/mpc/cheetah/arith/BUILD.bazel b/libspu/mpc/cheetah/arith/BUILD.bazel
index c862ba28..809f63cd 100644
--- a/libspu/mpc/cheetah/arith/BUILD.bazel
+++ b/libspu/mpc/cheetah/arith/BUILD.bazel
@@ -119,3 +119,9 @@ spu_cc_test(
         "//libspu/mpc/utils:simulate",
     ],
 )
+
+spu_cc_test(
+    name = "vector_encoder_test",
+    srcs = ["vector_encoder_test.cc"],
+    deps = [":arith_comm"],
+)
diff --git a/libspu/mpc/cheetah/arith/vector_encoder.cc b/libspu/mpc/cheetah/arith/vector_encoder.cc
index 51df186a..52b71488 100644
--- a/libspu/mpc/cheetah/arith/vector_encoder.cc
+++ b/libspu/mpc/cheetah/arith/vector_encoder.cc
@@ -92,7 +92,7 @@ void VectorEncoder::Backward(const NdArrayRef &vec, RLWEPt *out,
     xtmp[0] = xvec[0];
     // reverse and sign flip
     for (size_t i = 1; i < num_coeffs; ++i) {
-      xtmp[num_coeffs - 1 - i] = -xvec[i];
+      xtmp[poly_deg_ - i] = -xvec[i];
     }
 
     uint64_t *dst = out->data();


Test patch:
diff --git a/libspu/mpc/cheetah/arith/vector_encoder_test.cc b/libspu/mpc/cheetah/arith/vector_encoder_test.cc
new file mode 100644
index 00000000..6056a5c6
--- /dev/null
+++ b/libspu/mpc/cheetah/arith/vector_encoder_test.cc
@@ -0,0 +1,142 @@
+// Copyright 2024 Ant Group Co., Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "libspu/mpc/cheetah/arith/vector_encoder.h"
+
+#include "gtest/gtest.h"
+#include "seal/util/polyarithsmallmod.h"
+
+#include "libspu/mpc/cheetah/rlwe/modswitch_helper.h"
+#include "libspu/mpc/cheetah/rlwe/utils.h"
+#include "libspu/mpc/utils/ring_ops.h"
+
+namespace spu::mpc::cheetah::test {
+
+class VectorEncoderTest : public ::testing::TestWithParam<FieldType> {
+ protected:
+  static constexpr size_t poly_deg = 4096;
+
+  FieldType field_;
+  std::shared_ptr<ModulusSwitchHelper> ms_helper_;
+  std::shared_ptr<seal::SEALContext> context_;
+
+  inline uint32_t FieldBitLen(FieldType f) const { return 8 * SizeOf(f); }
+
+  void SetUp() override {
+    field_ = GetParam();
+    std::vector<int> modulus_bits;
+    switch (field_) {
+      case FieldType::FM32:
+        modulus_bits = {55, 39};
+        break;
+      case FieldType::FM64:
+        modulus_bits = {55, 55, 48};
+        break;
+      case FieldType::FM128:
+        modulus_bits = {59, 59, 59, 59, 50};
+        break;
+      default:
+        SPU_THROW("Not support field type {}", field_);
+    }
+
+    auto scheme_type = seal::scheme_type::ckks;
+    auto parms = seal::EncryptionParameters(scheme_type);
+    parms.set_poly_modulus_degree(poly_deg);
+    auto modulus = seal::CoeffModulus::Create(poly_deg, modulus_bits);
+    parms.set_coeff_modulus(modulus);
+    parms.set_use_special_prime(false);
+
+    context_ = std::make_shared<seal::SEALContext>(parms, true,
+                                                   seal::sec_level_type::none);
+    seal::SEALContext ms_context(parms, false, seal::sec_level_type::none);
+
+    uint32_t bitlen = FieldBitLen(field_);
+    ms_helper_ = std::make_shared<ModulusSwitchHelper>(ms_context, bitlen);
+  }
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    Cheetah, VectorEncoderTest,
+    testing::Values(FieldType::FM32, FieldType::FM64, FieldType::FM128),
+    [](const testing::TestParamInfo<VectorEncoderTest::ParamType> &p) {
+      return fmt::format("{}", p.param);
+    });
+
+bool DyadicProduct(RLWEPt &pt, const RLWEPt &oth,
+                   const seal::SEALContext &context) {
+  using namespace seal::util;
+  auto cntxt_data = context.get_context_data(pt.parms_id());
+  if (!cntxt_data) {
+    return false;
+  }
+
+  auto L = cntxt_data->parms().coeff_modulus().size();
+  if (pt.coeff_count() % L != 0) {
+    return false;
+  }
+
+  auto ntt_tables = cntxt_data->small_ntt_tables();
+  size_t n = pt.coeff_count() / L;
+  auto pt_ptr = pt.data();
+  auto oth_ptr = oth.data();
+  for (size_t l = 0; l < L; ++l) {
+    dyadic_product_coeffmod(pt_ptr, oth_ptr, n, ntt_tables[l].modulus(),
+                            pt_ptr);
+    pt_ptr += n;
+    oth_ptr += n;
+  }
+  return true;
+}
+
+TEST_P(VectorEncoderTest, ForwardBackward) {
+  VectorEncoder vencoder(*context_, *ms_helper_);
+  for (bool scale_forward : {true, false}) {
+    for (int64_t n : {100L, (long)poly_deg}) {
+      auto vec0 = ring_rand(field_, {n});
+      auto vec1 = ring_rand(field_, {n});
+
+      RLWEPt poly0;
+      RLWEPt poly1;
+      vencoder.Forward(vec0, &poly0, scale_forward);
+      vencoder.Backward(vec1, &poly1, !scale_forward);
+
+      NttInplace(poly0, *context_);
+      NttInplace(poly1, *context_);
+      DyadicProduct(poly0, poly1, *context_);
+      InvNttInplace(poly0, *context_);
+
+      size_t num_modulus = poly0.coeff_count() / poly_deg;
+      std::vector<uint64_t> cnst(num_modulus);
+      for (size_t l = 0; l < num_modulus; ++l) {
+        cnst[l] = poly0.data()[l * poly_deg];
+      }
+
+      auto computed =
+          ms_helper_->ModulusDownRNS(field_, {1L}, absl::MakeSpan(cnst));
+
+      DISPATCH_ALL_FIELDS(field_, "Check", [&]() {
+        NdArrayView<ring2k_t> got(computed);
+        NdArrayView<ring2k_t> v0(vec0);
+        NdArrayView<ring2k_t> v1(vec1);
+        ring2k_t expected = 0;
+        for (int64_t i = 0; i < v0.numel(); ++i) {
+          expected += v0[i] * v1[i];
+        }
+        ASSERT_EQ(expected, got[0]);
+      });
+    }
+  }
+}
+
+}  // namespace spu::mpc::cheetah::test
diff --git a/libspu/mpc/cheetah/rlwe/packlwes_test.cc b/libspu/mpc/cheetah/rlwe/packlwes_test.cc
index 6d8d8912..6d924c26 100644
--- a/libspu/mpc/cheetah/rlwe/packlwes_test.cc
+++ b/libspu/mpc/cheetah/rlwe/packlwes_test.cc
@@ -353,7 +353,7 @@ void VectorEncoder::Backward(const NdArrayRef &vec, RLWEPt *out,
     xtmp[0] = xvec[0];
     // reverse and sign flip
     for (size_t i = 1; i < num_coeffs; ++i) {
-      xtmp[num_coeffs - 1 - i] = -xvec[i];
+      xtmp[poly_deg_ - i] = -xvec[i];
     }
 
     uint64_t *dst = out->data();