Fix CKKS Multiply Bug (intel#98)

xiachonga · Jan 4, 2022 · 34e88ed · 34e88ed
1 parent ddd1d51
commit 34e88ed
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 8 deletions.
diff --git a/hexl/experimental/seal/ckks-multiply.cpp b/hexl/experimental/seal/ckks-multiply.cpp
@@ -54,16 +54,16 @@ void CkksMultiply(uint64_t* result, const uint64_t* operand1,
 
       // Compute second output polynomial
       // result[1] = x[1] * y[0]
-      intel::hexl::EltwiseMultMod(
-          &result[poly1_offset], operand1 + poly1_offset,
-          operand2 + poly0_offset, tile_size, moduli[i], 1);
-      // result[1] = x[0] * y[1]
-      intel::hexl::EltwiseMultMod(temp.data(), operand1 + poly0_offset,
-                                  operand2 + poly1_offset, tile_size, moduli[i],
+      intel::hexl::EltwiseMultMod(temp.data(), operand1 + poly1_offset,
+                                  operand2 + poly0_offset, tile_size, moduli[i],
                                   1);
+      // result[1] = x[0] * y[1]
+      intel::hexl::EltwiseMultMod(
+          &result[poly1_offset], operand1 + poly0_offset,
+          operand2 + poly1_offset, tile_size, moduli[i], 1);
       // result[1] += temp_poly
-      intel::hexl::EltwiseAddMod(&result[poly1_offset], &result[poly1_offset],
-                                 temp.data(), tile_size, moduli[i]);
+      intel::hexl::EltwiseAddMod(&result[poly1_offset], temp.data(),
+                                 &result[poly1_offset], tile_size, moduli[i]);
 
       // Compute first output polynomial
       // result[0] = x[0] * y[0]

diff --git a/test/experimental/seal/test-ckks-multiply.cpp b/test/experimental/seal/test-ckks-multiply.cpp
@@ -59,6 +59,58 @@ TEST(CkksMultiply, small_one_mod_inplace) {
   CheckEqual(op1, exp_out);
 }
 
+TEST(CkksMultiply, small_one_mod_square_same_op) {
+  size_t coeff_count = 3;
+  std::vector<uint64_t> moduli{10};
+
+  std::vector<uint64_t> op1{1, 2, 3,  //
+                            4, 5, 6};
+  std::vector<uint64_t> out(3 * coeff_count * moduli.size(), 0);
+
+  std::vector<uint64_t> exp_out{
+      (op1[0] * op1[0] % 10),
+      (op1[1] * op1[1] % 10),
+      (op1[2] * op1[2] % 10),  //
+      (op1[0] * op1[3] + op1[3] * op1[0]) % 10,
+      (op1[1] * op1[4] + op1[4] * op1[1]) % 10,
+      (op1[2] * op1[5] + op1[5] * op1[2]) % 10,  //
+      (op1[3] * op1[3] % 10),
+      (op1[4] * op1[4] % 10),
+      (op1[5] * op1[5] % 10)  //
+  };
+
+  CkksMultiply(out.data(), op1.data(), op1.data(), coeff_count, moduli.data(),
+               moduli.size());
+
+  CheckEqual(out, exp_out);
+}
+
+TEST(CkksMultiply, small_one_mod_square_same_op_inline) {
+  size_t coeff_count = 3;
+  std::vector<uint64_t> moduli{10};
+
+  std::vector<uint64_t> op1{1, 2, 3,  //
+                            4, 5, 6,  //
+                            0, 0, 0};
+
+  std::vector<uint64_t> exp_out{
+      (op1[0] * op1[0] % 10),
+      (op1[1] * op1[1] % 10),
+      (op1[2] * op1[2] % 10),  //
+      (op1[0] * op1[3] + op1[3] * op1[0]) % 10,
+      (op1[1] * op1[4] + op1[4] * op1[1]) % 10,
+      (op1[2] * op1[5] + op1[5] * op1[2]) % 10,  //
+      (op1[3] * op1[3] % 10),
+      (op1[4] * op1[4] % 10),
+      (op1[5] * op1[5] % 10)  //
+  };
+
+  CkksMultiply(op1.data(), op1.data(), op1.data(), coeff_count, moduli.data(),
+               moduli.size());
+
+  CheckEqual(op1, exp_out);
+}
+
 TEST(CkksMultiply, small_two_mod) {
   size_t coeff_count = 3;
   std::vector<uint64_t> moduli{10, 20};