radix sort support unsigned/signed type (#1367)

deadlywing · Copilot · web-flow · commit feb185aded76 · 2026-01-06T14:59:21.000+08:00
# Pull Request

## What problem does this PR solve?

Issue Number: Fixed #

## Possible side effects?

- Performance:

- Backward compatibility:

---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/src/libspu/core/type_util.cc b/src/libspu/core/type_util.cc
@@ -52,6 +52,14 @@ bool isInteger(DataType dtype) {
   }
 }
 
+bool isUInteger(DataType dtype) {
+  switch (dtype) {
+    FOREACH_UINT_DTYPES(CASE)
+    default:
+      return false;
+  }
+}
+
 bool isFixedPoint(DataType dtype) {
   switch (dtype) {
     FOREACH_FXP_DTYPES(CASE)
diff --git a/src/libspu/core/type_util.h b/src/libspu/core/type_util.h
@@ -50,6 +50,13 @@ std::ostream& operator<<(std::ostream& os, const Visibility& vtype);
   FN(DT_I64, I64, 64)          \
   FN(DT_U64, U64, 64)
 
+#define FOREACH_UINT_DTYPES(FN) \
+  FN(DT_I1, I1, 1)              \
+  FN(DT_U8, U8, 8)              \
+  FN(DT_U16, U16, 16)           \
+  FN(DT_U32, U32, 32)           \
+  FN(DT_U64, U64, 64)
+
 #define FOREACH_FXP_DTYPES(FN) \
   FN(DT_F16, F16, 16)          \
   FN(DT_F32, F32, 32)          \
@@ -60,6 +67,7 @@ std::ostream& operator<<(std::ostream& os, const Visibility& vtype);
   FOREACH_FXP_DTYPES(FN)
 
 bool isInteger(DataType dtype);
+bool isUInteger(DataType dtype);
 bool isFixedPoint(DataType dtype);
 size_t getWidth(DataType dtype);
 
diff --git a/src/libspu/core/value.h b/src/libspu/core/value.h
@@ -80,6 +80,7 @@ class Value final {
   // Get dtype.
   DataType dtype() const { return dtype_; }
   bool isInt() const { return isInteger(dtype()); }
+  bool isUInt() const { return isUInteger(dtype()); }
   bool isFxp() const { return isFixedPoint(dtype()); }
   bool isComplex() const { return imag_.has_value(); }
 
diff --git a/src/libspu/kernel/hal/permute.cc b/src/libspu/kernel/hal/permute.cc
@@ -984,25 +984,27 @@ std::vector<spu::Value> _gen_bv_vector(SPUContext *ctx,
                                        int64_t valid_bits) {
   std::vector<spu::Value> ret;
   const auto k1 = _constant(ctx, 1U, keys[0].shape());
+  const bool is_descending = (direction == SortDirection::Descending);
+
   // keys[0] is the most significant key
   for (size_t i = keys.size(); i > 0; --i) {
     const auto t = _bit_decompose(ctx, keys[i - 1], valid_bits);
+    const bool is_unsigned = keys[i - 1].isUInt();
 
     SPU_ENFORCE(t.size() > 0);
+
+    // Process non-sign bits (same logic for both signed and unsigned types)
+    // Radix sort is a stable sorting algorithm for the ascending order, if
+    // we flip the bit, then we can get the descending order for stable sort
     for (size_t j = 0; j < t.size() - 1; j++) {
-      // Radix sort is a stable sorting algorithm for the ascending order, if
-      // we flip the bit, then we can get the descending order for stable sort
-      if (direction == SortDirection::Descending) {
-        ret.emplace_back(_sub(ctx, k1, t[j]));
-      } else {
-        ret.emplace_back(t[j]);
-      }
+      ret.emplace_back(is_descending ? _sub(ctx, k1, t[j]) : t[j]);
     }
-    // The sign bit is opposite
-    if (direction == SortDirection::Descending) {
-      ret.emplace_back(t.back());
+
+    // For signed types, handle the sign bit (opposite logic)
+    if (!is_unsigned) {
+      ret.emplace_back(is_descending ? t.back() : _sub(ctx, k1, t.back()));
     } else {
-      ret.emplace_back(_sub(ctx, k1, t.back()));
+      ret.emplace_back(is_descending ? _sub(ctx, k1, t.back()) : t.back());
     }
   }
   return ret;
diff --git a/src/libspu/kernel/hal/permute.h b/src/libspu/kernel/hal/permute.h
@@ -58,6 +58,10 @@ std::vector<spu::Value> sort1d(SPUContext *ctx,
 //  - direction: sorting order
 //  - num_keys: the number of operands to treat as keys (count from index 0)
 //  - valid_bits: indicates the numeric range of keys for performance hint
+//
+// Important notes:
+//   - for radix sort, the user should ensure that the data has the correct
+//     signed or unsigned type.
 std::vector<spu::Value> simple_sort1d(SPUContext *ctx,
                                       absl::Span<spu::Value const> inputs,
                                       SortDirection direction, int64_t num_keys,
diff --git a/src/libspu/kernel/hlo/sort_test.cc b/src/libspu/kernel/hlo/sort_test.cc
@@ -15,6 +15,7 @@
 #include "libspu/kernel/hlo/sort.h"
 
 #include <algorithm>
+#include <limits>
 #include <random>
 #include <xtensor/xadapt.hpp>
 #include <xtensor/xsort.hpp>
@@ -461,6 +462,244 @@ TEST_P(SimpleSortTest, MixVisibilityKey) {
       });
 }
 
+// Helper template to run unsigned type sort test
+template <typename T>
+void RunUnsignedSortTest(SPUContext *ctx) {
+  xt::xarray<T> k1 = {7, 6, 5, 4, 1, 3, 2};
+  xt::xarray<T> sorted_k1 = {1, 2, 3, 4, 5, 6, 7};
+  xt::xarray<float> payload = {1, 2, 3, 6, 7, 6, 5};
+  xt::xarray<float> sorted_payload = {7, 5, 6, 6, 3, 2, 1};
+
+  Value k1_v = test::makeValue(ctx, k1, VIS_SECRET);
+  Value payload_v = test::makeValue(ctx, payload, VIS_SECRET);
+
+  std::vector<spu::Value> rets =
+      SimpleSort(ctx, {k1_v, payload_v}, 0, hal::SortDirection::Ascending, 1);
+
+  EXPECT_EQ(rets.size(), 2);
+
+  auto sorted_k1_hat = hal::dump_public_as<T>(ctx, hal::reveal(ctx, rets[0]));
+  auto sorted_payload_hat =
+      hal::dump_public_as<float>(ctx, hal::reveal(ctx, rets[1]));
+
+  EXPECT_TRUE(xt::allclose(sorted_k1, sorted_k1_hat, 0.01, 0.001))
+      << "sort failed: " << sorted_k1 << std::endl
+      << sorted_k1_hat << std::endl;
+
+  EXPECT_TRUE(xt::allclose(sorted_payload, sorted_payload_hat, 0.01, 0.001))
+      << "payload failed: " << sorted_payload << std::endl
+      << sorted_payload_hat << std::endl;
+}
+
+TEST_P(SimpleSortTest, UnsignedTypeSort) {
+  size_t npc = std::get<0>(GetParam());
+  FieldType field = std::get<1>(GetParam());
+  ProtocolKind prot = std::get<2>(GetParam());
+  RuntimeConfig::SortMethod method = std::get<3>(GetParam());
+
+  mpc::utils::simulate(npc,
+                       [&](const std::shared_ptr<yacl::link::Context> &lctx) {
+                         RuntimeConfig cfg;
+                         cfg.protocol = prot;
+                         cfg.field = field;
+                         cfg.enable_action_trace = false;
+                         cfg.sort_method = method;
+                         SPUContext ctx = test::makeSPUContext(cfg, lctx);
+
+                         RunUnsignedSortTest<uint8_t>(&ctx);
+                         RunUnsignedSortTest<uint16_t>(&ctx);
+                         RunUnsignedSortTest<uint32_t>(&ctx);
+                         if (field >= FieldType::FM64) {
+                           RunUnsignedSortTest<uint64_t>(&ctx);
+                         }
+                       });
+}
+
+// Helper template to test signed interpretation of unsigned-range values
+// When values like 255 (for int8_t) are interpreted as signed, they become -1
+// So sorting should treat them as negative numbers
+template <typename SignedT, typename UnsignedT>
+void RunSignedInterpretationSortTest(SPUContext *ctx) {
+  // Use max value of unsigned type which becomes -1 when interpreted as signed
+  constexpr UnsignedT max_val = std::numeric_limits<UnsignedT>::max();
+  // Key: {0, max_val} where max_val is interpreted as -1 in signed
+  // For ascending sort with signed interpretation: -1 < 0, so max_val comes
+  // first
+  xt::xarray<SignedT> k1 = {0, static_cast<SignedT>(max_val)};
+  // Expected: max_val (-1) < 0, so sorted order is {max_val, 0}
+  xt::xarray<SignedT> sorted_k1 = {static_cast<SignedT>(max_val), 0};
+  xt::xarray<float> payload = {1.0, 2.0};
+  xt::xarray<float> sorted_payload = {2.0, 1.0};
+
+  Value k1_v = test::makeValue(ctx, k1, VIS_SECRET);
+  Value payload_v = test::makeValue(ctx, payload, VIS_SECRET);
+
+  std::vector<spu::Value> rets =
+      SimpleSort(ctx, {k1_v, payload_v}, 0, hal::SortDirection::Ascending, 1);
+
+  EXPECT_EQ(rets.size(), 2);
+
+  auto sorted_k1_hat =
+      hal::dump_public_as<SignedT>(ctx, hal::reveal(ctx, rets[0]));
+  auto sorted_payload_hat =
+      hal::dump_public_as<float>(ctx, hal::reveal(ctx, rets[1]));
+
+  EXPECT_TRUE(xt::allclose(sorted_k1, sorted_k1_hat, 0.01, 0.001))
+      << "sort failed: expected " << sorted_k1 << ", got " << sorted_k1_hat
+      << std::endl;
+
+  EXPECT_TRUE(xt::allclose(sorted_payload, sorted_payload_hat, 0.01, 0.001))
+      << "payload failed: expected " << sorted_payload << ", got "
+      << sorted_payload_hat << std::endl;
+}
+
+// IMPORTANT: the user should ensure that the data has the correct signed or
+// unsigned type. Incorrect type interpretation will result in incorrect sort
+// order (for example, treating signed values as unsigned may place negative
+// numbers at the end instead of the beginning).
+TEST_P(SimpleSortTest, SignedInterpretationSort) {
+  size_t npc = std::get<0>(GetParam());
+  FieldType field = std::get<1>(GetParam());
+  ProtocolKind prot = std::get<2>(GetParam());
+  RuntimeConfig::SortMethod method = std::get<3>(GetParam());
+
+  mpc::utils::simulate(
+      npc, [&](const std::shared_ptr<yacl::link::Context> &lctx) {
+        RuntimeConfig cfg;
+        cfg.protocol = prot;
+        cfg.field = field;
+        cfg.enable_action_trace = false;
+        cfg.sort_method = method;
+        SPUContext ctx = test::makeSPUContext(cfg, lctx);
+
+        // Test: data is uint8_t range but treated as int8_t
+        // 255 (uint8_t) -> -1 (int8_t), so -1 < 0
+        RunSignedInterpretationSortTest<int8_t, uint8_t>(&ctx);
+
+        // Test: data is uint16_t range but treated as int16_t
+        // 65535 (uint16_t) -> -1 (int16_t), so -1 < 0
+        RunSignedInterpretationSortTest<int16_t, uint16_t>(&ctx);
+
+        // Test: data is uint32_t range but treated as int32_t
+        // 4294967295 (uint32_t) -> -1 (int32_t), so -1 < 0
+        RunSignedInterpretationSortTest<int32_t, uint32_t>(&ctx);
+
+        if (field >= FieldType::FM64) {
+          // Test: data is uint64_t range but treated as int64_t
+          RunSignedInterpretationSortTest<int64_t, uint64_t>(&ctx);
+        }
+      });
+}
+
+TEST_P(SimpleSortTest, BoolKeyWithPayloads) {
+  size_t npc = std::get<0>(GetParam());
+  FieldType field = std::get<1>(GetParam());
+  ProtocolKind prot = std::get<2>(GetParam());
+  RuntimeConfig::SortMethod method = std::get<3>(GetParam());
+
+  mpc::utils::simulate(
+      npc, [&](const std::shared_ptr<yacl::link::Context> &lctx) {
+        RuntimeConfig cfg;
+        cfg.protocol = prot;
+        cfg.field = field;
+        cfg.enable_action_trace = false;
+        cfg.sort_method = method;
+
+        SPUContext ctx = test::makeSPUContext(cfg, lctx);
+
+        // Bool key with two payloads
+        xt::xarray<bool> k1 = {true, false, true, false, true};
+        xt::xarray<float> p1 = {1.0, 2.0, 3.0, 4.0, 5.0};
+        xt::xarray<int32_t> p2 = {10, 20, 30, 40, 50};
+
+        // Expected sorted keys
+        xt::xarray<bool> sorted_k1_desc = {true, true, true, false, false};
+        xt::xarray<bool> sorted_k1_asc = {false, false, true, true, true};
+
+        // Expected payloads (sorted within each group since sort is unstable)
+        // Descending: true keys first {1,3,5}, then false keys {2,4}
+        xt::xarray<float> sorted_p1_desc = {1.0, 3.0, 5.0, 2.0, 4.0};
+        xt::xarray<int32_t> sorted_p2_desc = {10, 30, 50, 20, 40};
+        // Ascending: false keys first {2,4}, then true keys {1,3,5}
+        xt::xarray<float> sorted_p1_asc = {2.0, 4.0, 1.0, 3.0, 5.0};
+        xt::xarray<int32_t> sorted_p2_asc = {20, 40, 10, 30, 50};
+
+        Value k1_v = test::makeValue(&ctx, k1, VIS_SECRET);
+        Value p1_v = test::makeValue(&ctx, p1, VIS_SECRET);
+        Value p2_v = test::makeValue(&ctx, p2, VIS_SECRET);
+
+        // Test descending sort (true before false)
+        {
+          std::vector<spu::Value> rets = SimpleSort(
+              &ctx, {k1_v, p1_v, p2_v}, 0, hal::SortDirection::Descending, 1);
+
+          EXPECT_EQ(rets.size(), 3);
+
+          auto sorted_k1_hat =
+              hal::dump_public_as<bool>(&ctx, hal::reveal(&ctx, rets[0]));
+          auto sorted_p1_hat =
+              hal::dump_public_as<float>(&ctx, hal::reveal(&ctx, rets[1]));
+          auto sorted_p2_hat =
+              hal::dump_public_as<int32_t>(&ctx, hal::reveal(&ctx, rets[2]));
+
+          // Check bool key is sorted correctly
+          EXPECT_TRUE(xt::allclose(sorted_k1_desc, sorted_k1_hat, 0.01, 0.001))
+              << "Bool descending sort failed: " << sorted_k1_desc << std::endl
+              << sorted_k1_hat << std::endl;
+
+          // Sort each part and compare (since sort is unstable within same key)
+          auto p1_hat_sorted = xt::concatenate(
+              xt::xtuple(xt::sort(xt::view(sorted_p1_hat, xt::range(0, 3))),
+                         xt::sort(xt::view(sorted_p1_hat, xt::range(3, 5)))));
+          auto p2_hat_sorted = xt::concatenate(
+              xt::xtuple(xt::sort(xt::view(sorted_p2_hat, xt::range(0, 3))),
+                         xt::sort(xt::view(sorted_p2_hat, xt::range(3, 5)))));
+
+          EXPECT_TRUE(xt::allclose(sorted_p1_desc, p1_hat_sorted, 0.01, 0.001))
+              << "Descending p1 failed: " << sorted_p1_desc << std::endl
+              << p1_hat_sorted << std::endl;
+          EXPECT_TRUE(xt::allclose(sorted_p2_desc, p2_hat_sorted, 0.01, 0.001))
+              << "Descending p2 failed: " << sorted_p2_desc << std::endl
+              << p2_hat_sorted << std::endl;
+        }
+
+        // Test ascending sort (false before true)
+        {
+          std::vector<spu::Value> rets = SimpleSort(
+              &ctx, {k1_v, p1_v, p2_v}, 0, hal::SortDirection::Ascending, 1);
+
+          EXPECT_EQ(rets.size(), 3);
+
+          auto sorted_k1_hat =
+              hal::dump_public_as<bool>(&ctx, hal::reveal(&ctx, rets[0]));
+          auto sorted_p1_hat =
+              hal::dump_public_as<float>(&ctx, hal::reveal(&ctx, rets[1]));
+          auto sorted_p2_hat =
+              hal::dump_public_as<int32_t>(&ctx, hal::reveal(&ctx, rets[2]));
+
+          // Check bool key is sorted correctly
+          EXPECT_TRUE(xt::allclose(sorted_k1_asc, sorted_k1_hat, 0.01, 0.001))
+              << "Bool ascending sort failed: " << sorted_k1_asc << std::endl
+              << sorted_k1_hat << std::endl;
+
+          // Sort each part and compare (since sort is unstable within same key)
+          auto p1_hat_sorted = xt::concatenate(
+              xt::xtuple(xt::sort(xt::view(sorted_p1_hat, xt::range(0, 2))),
+                         xt::sort(xt::view(sorted_p1_hat, xt::range(2, 5)))));
+          auto p2_hat_sorted = xt::concatenate(
+              xt::xtuple(xt::sort(xt::view(sorted_p2_hat, xt::range(0, 2))),
+                         xt::sort(xt::view(sorted_p2_hat, xt::range(2, 5)))));
+
+          EXPECT_TRUE(xt::allclose(sorted_p1_asc, p1_hat_sorted, 0.01, 0.001))
+              << "Ascending p1 failed: " << sorted_p1_asc << std::endl
+              << p1_hat_sorted << std::endl;
+          EXPECT_TRUE(xt::allclose(sorted_p2_asc, p2_hat_sorted, 0.01, 0.001))
+              << "Ascending p2 failed: " << sorted_p2_asc << std::endl
+              << p2_hat_sorted << std::endl;
+        }
+      });
+}
+
 INSTANTIATE_TEST_SUITE_P(
     SimpleSort2PCTestInstances, SimpleSortTest,
     testing::Combine(testing::Values(2),

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,14 @@ bool isInteger(DataType dtype) {`
`52`	`52`	`}`
`53`	`53`	`}`
`54`	`54`
	`55`	`+bool isUInteger(DataType dtype) {`
	`56`	`+ switch (dtype) {`
	`57`	`+ FOREACH_UINT_DTYPES(CASE)`
	`58`	`+ default:`
	`59`	`+ return false;`
	`60`	`+ }`
	`61`	`+}`
	`62`	`+`
`55`	`63`	`bool isFixedPoint(DataType dtype) {`
`56`	`64`	`switch (dtype) {`
`57`	`65`	`FOREACH_FXP_DTYPES(CASE)`