[squirrel] Remove the assertation on empty indicator, only reporting …

experimental/squirrel/README.md

@@ -29,11 +29,24 @@ Code under this folder is purely for research demonstration and it's **NOT desig
   * On one terminal
 
     ```sh
-    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=85 --rank1_nfeatures=85 --standalone --train=BinaryClassification_Aps_Test_60000_171.csv --test=BinaryClassification_Aps_Test_16000_171.csv --standalone=true --rank=0 --has_label=0 --lr=1.0 --subsample=0.8
+    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=85 --rank1_nfeatures=85 --standalone=true --train=BinaryClassification_Aps_Test_60000_171.csv --test=BinaryClassification_Aps_Test_16000_171.csv --rank=0 --has_label=0 --lr=1.0 --subsample=0.8
     ```
 
   * On another terminal
 
     ```sh
-    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=85 --rank1_nfeatures=85 --standalone --train=BinaryClassification_Aps_Test_60000_171.csv --test=BinaryClassification_Aps_Test_16000_171.csv --standalone=true --rank=1 --has_label=1 --lr=1.0 --subsample=0.8
+    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=85 --rank1_nfeatures=85 --standalone=true --train=BinaryClassification_Aps_Test_60000_171.csv --test=BinaryClassification_Aps_Test_16000_171.csv --rank=1 --has_label=1 --lr=1.0 --subsample=0.8
     ```
+* Run on distributed dataset, e.g., using the `breast_cancer` dataset from the SPU repo. 
+  * On one terminal
+
+    ```sh
+    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=15 --rank1_nfeatures=15 --standalone=false --train=examples/data/breast_cancer_a.csv --rank=0 --has_label=0 --lr=1.0 --subsample=0.8
+    ```
+
+  * On another terminal
+
+    ```sh
+    bazel-bin/experimental/squirrel/squirrel_demo_main --rank0_nfeatures=15 --rank1_nfeatures=15 --standalone=false --train=examples/data/breast_cancer_b.csv --rank=1 --has_label=1 --lr=1.0 --subsample=0.8
+    ```
+

experimental/squirrel/bin_matvec_prot.cc

@@ -766,9 +766,12 @@ spu::NdArrayRef BinMatVecProtocol::Recv(const spu::NdArrayRef &vec_in,
   SPU_ENFORCE_EQ(vec_in.numel(), dim_in);
   if (not indicator.empty()) {
     // mat * diag(indicator) should not be all zeros.
-    SPU_ENFORCE(std::any_of(indicator.begin(), indicator.end(),
-                            [](uint8_t x) { return x > 0; }),
-                "empty matrix is not allowed");
+    if (std::all_of(indicator.begin(), indicator.end(),
+                    [](uint8_t x) { return x == 0; })) {
+      SPDLOG_WARN(
+          "Empty matrix! Make sure the 1-bit error will not ruin your "
+          "computation.");
+    }
   }
 
   auto eltype = vec_in.eltype();

experimental/squirrel/bin_matvec_prot.h

@@ -70,7 +70,10 @@ struct StlSparseMatrix {
 // Outputs:
 //   Sender: z0 \in Zk^{m}
 //     Recv: z1 \in Zk^{m}
-// such that z0 + z1 = M * (v0 + v1) mod Zk
+// such that z0 + z1 = M * (v0 + v1) + e mod Zk
+//
+// Note that we might introduce 1-bit error `e` due to the coefficient-based
+// resharing HE ciphertexts to additive shares.
 class BinMatVecProtocol {
  public:
   BinMatVecProtocol(size_t ring_bitwidth,

experimental/squirrel/bin_matvec_prot_test.cc

@@ -172,4 +172,100 @@ TEST_P(BinMatVecProtTest, WithIndicator) {
     }
   });
 }
+
+TEST_P(BinMatVecProtTest, EmptyMat) {
+  using namespace spu;
+  using namespace spu::mpc;
+  constexpr size_t kWorldSize = 2;
+
+  FieldType field = std::get<0>(GetParam());
+  int64_t dim_in = std::get<0>(std::get<1>(GetParam()));
+  int64_t dim_out = std::get<1>(std::get<1>(GetParam()));
+
+  StlSparseMatrix mat;
+  mat.rows_data_.resize(dim_out);
+  mat.cols_ = dim_in;
+
+  NdArrayRef vec_shr[2];
+  vec_shr[0] = ring_rand(field, {dim_in})
+                   .as(spu::makeType<spu::mpc::cheetah::AShrTy>(field));
+  vec_shr[1] = ring_rand(field, {dim_in})
+                   .as(spu::makeType<spu::mpc::cheetah::AShrTy>(field));
+
+  NdArrayRef vec = ring_add(vec_shr[0], vec_shr[1]);
+
+  NdArrayRef out_shr[2];
+  utils::simulate(kWorldSize, [&](std::shared_ptr<yacl::link::Context> lctx) {
+    BinMatVecProtocol binmat_prot(SizeOf(field) * 8, lctx);
+    if (0 == lctx->Rank()) {
+      out_shr[0] = binmat_prot.Send(vec_shr[0], dim_out, dim_in);
+    } else {
+      out_shr[1] = binmat_prot.Recv(vec_shr[1], dim_out, dim_in, mat);
+    }
+  });
+  NdArrayRef reveal = ring_add(out_shr[0], out_shr[1]);
+
+  DISPATCH_ALL_FIELDS(field, "", [&]() {
+    using sT = std::make_signed<ring2k_t>::type;
+    NdArrayView<sT> _vec(vec);
+    auto expected = BinAccumuate<sT>(_vec, mat);
+    NdArrayView<sT> got(reveal);
+
+    EXPECT_EQ(expected.size(), (size_t)got.numel());
+    for (int64_t i = 0; i < dim_out; ++i) {
+      EXPECT_NEAR(expected[i], got[i], 1);
+    }
+  });
+}
+
+TEST_P(BinMatVecProtTest, WithEmptyIndicator) {
+  using namespace spu;
+  using namespace spu::mpc;
+  constexpr size_t kWorldSize = 2;
+
+  FieldType field = std::get<0>(GetParam());
+  int64_t dim_in = std::get<0>(std::get<1>(GetParam()));
+  int64_t dim_out = std::get<1>(std::get<1>(GetParam()));
+
+  StlSparseMatrix mat;
+  PrepareBinaryMat(mat, dim_out, dim_in, 0);
+
+  NdArrayRef vec_shr[2];
+  vec_shr[0] = ring_rand(field, {dim_in})
+                   .as(spu::makeType<spu::mpc::cheetah::AShrTy>(field));
+  vec_shr[1] = ring_rand(field, {dim_in})
+                   .as(spu::makeType<spu::mpc::cheetah::AShrTy>(field));
+  std::vector<uint8_t> indicator(dim_in);
+  std::default_random_engine rdv;
+  std::uniform_int_distribution<uint8_t> dist(0, 10);
+  // empty indicator
+  std::fill_n(indicator.data(), indicator.size(), 0);
+
+  NdArrayRef vec = ring_add(vec_shr[0], vec_shr[1]);
+
+  NdArrayRef out_shr[2];
+  utils::simulate(kWorldSize, [&](std::shared_ptr<yacl::link::Context> lctx) {
+    BinMatVecProtocol binmat_prot(SizeOf(field) * 8, lctx);
+    if (0 == lctx->Rank()) {
+      out_shr[0] = binmat_prot.Send(vec_shr[0], dim_out, dim_in);
+    } else {
+      out_shr[1] = binmat_prot.Recv(vec_shr[1], dim_out, dim_in, mat,
+                                    absl::MakeConstSpan(indicator));
+    }
+  });
+  NdArrayRef reveal = ring_add(out_shr[0], out_shr[1]);
+
+  DISPATCH_ALL_FIELDS(field, "", [&]() {
+    using sT = std::make_signed<ring2k_t>::type;
+    NdArrayView<sT> _vec(vec);
+    NdArrayView<sT> got(reveal);
+    auto expected = BinAccumuate<sT>(_vec, mat, absl::MakeConstSpan(indicator));
+
+    EXPECT_EQ(expected.size(), (size_t)got.numel());
+    for (int64_t i = 0; i < dim_out; ++i) {
+      EXPECT_NEAR(expected[i], got[i], 1);
+    }
+  });
+}
+
 }  // namespace squirrel::test