repo-sync-2025-03-19T11:04:50+0800 (#1074)

Author: FollyCoolly

MODULE.bazel

@@ -21,7 +21,7 @@
 
 module(
     name = "spu",
-    version = "0.9.4.dev20250312",
+    version = "0.9.4.dev20250319",
     compatibility_level = 1,
 )
 

src/MODULE.bazel

@@ -21,7 +21,7 @@
 
 module(
     name = "spulib",
-    version = "0.9.4.dev20250312",
+    version = "0.9.4.dev20250319",
     compatibility_level = 1,
 )
 

src/libspu/kernel/hal/permute.cc

@@ -611,8 +611,8 @@ std::vector<spu::Value> PrepareSort(SPUContext *ctx,
   // use a random permutation to break link of values, such that the following
   // comparison can be revealed without loss of information.
   for (const auto &input : inputs) {
-    inp.emplace_back(
-        std::move(_perm_ss(ctx, input, rand_perm).setDtype(input.dtype())));
+    inp.emplace_back(std::move(
+        _perm_ss(ctx, _2s(ctx, input), rand_perm).setDtype(input.dtype())));
   }
 
   return inp;
@@ -1608,6 +1608,12 @@ std::vector<spu::Value> simple_sort1d(SPUContext *ctx,
   //   and the number of rounds increases (poly) logarithmically. In contrast,
   //   when the ring size doubles in radix sort, the communication （roughly）
   //   quadruples and the number of rounds doubles.
+  //   6. The above conclusions regarding performance apply only to
+  //   the cases of SECRET input and SECRET permutation. In reality, only radix
+  //   sort has implemented a complete mechanism for selecting the best
+  //   implementation based on visibility. The other implementations will use
+  //   local computation only when all keys are public; in other cases, they
+  //   will revert to the scenarios of SECRET input and SECRET permutation.
   //
 
   // if all keys are public, fallback to plaintext sort.

src/libspu/kernel/hal/soprf.cc

@@ -23,6 +23,10 @@ namespace spu::kernel::hal {
 Value soprf(SPUContext* ctx, const Value& x) {
   SPU_TRACE_HAL_LEAF(ctx, x);
 
+  if (x.numel() == 0) {
+    return x;
+  }
+
   // currently, wo only support LowMC block cipher
   SPU_ENFORCE(ctx->hasKernel("lowmc_b"));
   auto inp = x;
@@ -64,6 +68,10 @@ Value soprf(SPUContext* ctx, absl::Span<const spu::Value> inputs) {
                           }),
               "not all element has same dtype");
 
+  if (inputs.front().numel() == 0) {
+    return inputs.front();
+  }
+
   std::vector<Value> inp;
   inp.reserve(inputs.size());
   for (const auto& v : inputs) {

src/libspu/kernel/hlo/BUILD.bazel

@@ -349,6 +349,9 @@ spu_cc_test(
     deps = [
         ":shuffle",
         "//libspu/kernel:test_util",
+        "//libspu/kernel/hlo:casting",
+        "//libspu/kernel/hlo:const",
+        "//libspu/mpc/utils:simulate",
     ],
 )
 

src/libspu/kernel/hlo/permute_test.cc

@@ -123,8 +123,6 @@ xt::xarray<T> evalSinglePermuteOp(SPUContext* ctx, VisType x_vis,
                                   PtBufferView perm,
                                   const PermuteFunc& perm_func,
                                   int64_t perm_dim = 0) {
-  const auto prot = ctx->config().protocol;
-
   auto x_v = makeTestValue(ctx, x, x_vis);
   auto perm_v = makeTestValue(ctx, perm, perm_vis);
 
@@ -139,9 +137,7 @@ xt::xarray<T> evalSinglePermuteOp(SPUContext* ctx, VisType x_vis,
     EXPECT_EQ(send_round, 0);
   }
 
-  // costs of cheetah is highly dependant of OT kind, so we skip it.
-  if (prot != CHEETAH && ctx->hasKernel("inv_perm_av") &&
-      checkSpPass(x_vis, perm_vis)) {
+  if (ctx->hasKernel("inv_perm_av") && checkSpPass(x_vis, perm_vis)) {
     auto n_repeat = x_v.shape().numel() / x_v.shape().dim(perm_dim);
     // For ss version, at least 3 rounds.
     EXPECT_LE(std::min(send_round, recv_round), 2 * n_repeat);
@@ -194,11 +190,10 @@ std::vector<PermuteParams> GetValidParamsCombinations() {
 
   for (const auto& vis_x : kVisTypes) {
     for (const auto& vis_perm : kVisTypes) {
-      for (const auto& protocol : {CHEETAH, SEMI2K, ABY3}) {
+      for (const auto& protocol : {SEMI2K, ABY3}) {
         for (const auto& npc : {2, 3}) {
-          // npc=2/3 is not valid in ABY3/CHEETAH
-          if ((protocol == ABY3 && npc == 2) ||
-              (protocol == CHEETAH && npc == 3)) {
+          // npc=2 is not valid in ABY3
+          if (protocol == ABY3 && npc == 2) {
             continue;  // Skip invalid combinations
           }
           valid_combinations.emplace_back(vis_x, vis_perm, protocol, npc);
@@ -327,17 +322,14 @@ TEST_P(PermuteTest, MultiplePermuteWork) {
 class PermuteEmptyTest : public ::testing::TestWithParam<ProtocolKind> {};
 
 INSTANTIATE_TEST_SUITE_P(
-    PermuteEmpty, PermuteEmptyTest, testing::Values(CHEETAH, SEMI2K, ABY3),
+    PermuteEmpty, PermuteEmptyTest, testing::Values(SEMI2K, ABY3),
     [](const testing::TestParamInfo<PermuteEmptyTest::ParamType>& p) {
       return fmt::format("{}", p.param);
     });
 
 TEST_P(PermuteEmptyTest, Empty) {
   ProtocolKind prot = GetParam();
   size_t npc = 3;
-  if (prot == CHEETAH) {
-    npc = 2;
-  }
 
   mpc::utils::simulate(
       npc, [&](const std::shared_ptr<yacl::link::Context>& lctx) {

src/libspu/kernel/hlo/rank.cc

@@ -79,6 +79,14 @@ std::vector<spu::Value> TopK(SPUContext *ctx, const spu::Value &input,
                              int64_t k_lo, int64_t k_hi, bool largest,
                              bool value_only) {
   const Shape &shape = input.shape();
+
+  if (shape.numel() == 0) {
+    if (value_only) {
+      return std::vector<spu::Value>(1, input);
+    }
+    return std::vector<spu::Value>(2, input);
+  }
+
   SPU_ENFORCE(shape.numel() > 0, "input must non-empty.");
   SPU_ENFORCE(
       k_lo <= shape.back() && k_lo > 0,

src/libspu/kernel/hlo/rank_test.cc

@@ -374,12 +374,32 @@ TEST_P(TopkTest, ValueOnlyTest) {
       });
 }
 
+TEST_P(TopkTest, EmptyTest) {
+  size_t npc = std::get<0>(GetParam());
+  FieldType field = std::get<1>(GetParam());
+  ProtocolKind prot = std::get<2>(GetParam());
+
+  mpc::utils::simulate(
+      npc, [&](const std::shared_ptr<yacl::link::Context> &lctx) {
+        SPUContext sctx = test::makeSPUContext(prot, field, lctx);
+        auto empty_x = test::makeValue(&sctx, 1, VIS_SECRET, DT_INVALID, {0});
+
+        auto out = TopK(&sctx, empty_x, 1, 1);
+        EXPECT_EQ(out.size(), 2);
+        EXPECT_EQ(out[0].numel(), 0);
+        EXPECT_EQ(out[1].numel(), 0);
+        EXPECT_EQ(out[0].shape().size(), 1);
+        EXPECT_EQ(out[1].shape().size(), 1);
+        EXPECT_EQ(out[0].shape()[0], 0);
+        EXPECT_EQ(out[1].shape()[0], 0);
+      });
+}
+
 INSTANTIATE_TEST_SUITE_P(
     Topk2PCTestInstances, TopkTest,
     testing::Combine(testing::Values(2),
                      testing::Values(FieldType::FM64, FieldType::FM128),
-                     testing::Values(ProtocolKind::SEMI2K,
-                                     ProtocolKind::CHEETAH)),
+                     testing::Values(ProtocolKind::SEMI2K)),
     [](const testing::TestParamInfo<TopkTest::ParamType> &p) {
       return fmt::format("{}x{}x{}", std::get<0>(p.param), std::get<1>(p.param),
                          std::get<2>(p.param));

src/libspu/kernel/hlo/shuffle.cc

@@ -39,19 +39,26 @@ std::vector<spu::Value> Shuffle(SPUContext* ctx,
                                 absl::Span<const spu::Value> inputs,
                                 int64_t axis) {
   SPU_ENFORCE_GT(inputs.size(), 0U);
-  if (inputs[0].numel() == 0) {
+  auto input_shape = inputs[0].shape();
+  SPU_ENFORCE(std::all_of(inputs.begin() + 1, inputs.end(),
+                          [&](const spu::Value& v) {
+                            return v.shape() == input_shape;
+                          }),
+              "all inputs should have the same shape");
+
+  // edge case: empty or single element tensor
+  if (inputs[0].numel() <= 1) {
     return std::vector<spu::Value>(inputs.begin(), inputs.end());
   }
-  auto input_shape = inputs[0].shape();
 
   // TODO: Rename permute-related kernels
   if (ctx->hasKernel("rand_perm_m") && ctx->hasKernel("perm_am")) {
     auto shuffle_fn = [&](absl::Span<const spu::Value> input) {
       std::vector<spu::Value> rets;
-      auto rand_perm = hal::_rand_perm_s(ctx, input_shape);
-      for (size_t i = 0; i < input.size(); ++i) {
-        rets.emplace_back(hal::_perm_ss(ctx, _2s(ctx, input[i]), rand_perm)
-                              .setDtype(input[i].dtype()));
+      auto rand_perm = hal::_rand_perm_s(ctx, {input_shape.dim(axis)});
+      for (const auto& inp : input) {
+        rets.emplace_back(
+            hal::_perm_ss(ctx, _2s(ctx, inp), rand_perm).setDtype(inp.dtype()));
       }
       return rets;
     };