Merge pull request #3544 from LittlePyx:fix/vqsort-network-signed-zeros

PiperOrigin-RevId: 1019162426

Author: copybara-github

hwy/contrib/sort/shared-inl.h

@@ -165,6 +165,10 @@ static_assert(SortConstants::MaxBufBytes<2>(64) <= 1664, "Unexpectedly high");
 namespace hwy {
 namespace HWY_NAMESPACE {
 
+namespace detail {
+enum class PairLayout { kIndependent, kMirrored };
+}  // namespace detail
+
 // Default tag / vector width selector.
 #if HWY_TARGET == HWY_RVV
 // Use LMUL = 1/2; for SEW=64 this ends up emulated via VSETVLI.

hwy/contrib/sort/sort_unit_test.cc

@@ -309,6 +309,27 @@ static HWY_NOINLINE void TestBaseCase() {
   TestBaseCase01<Traits>();
 }
 
+template <class Traits>
+static HWY_NOINLINE void TestBaseCaseSignedZero() {
+  using T = typename Traits::LaneType;
+  const SortTag<T> d;
+  const size_t num = SortConstants::BaseCaseNumLanes(Lanes(d));
+  auto keys = hwy::AllocateAligned<T>(num);
+  auto buf = hwy::AllocateAligned<T>(SortConstants::BufNum<T>(Lanes(d)));
+  HWY_ASSERT(keys && buf);
+  // Two keys exercise one compare-exchange. The full network also compares
+  // lanes within the same vector; place the negative zero on either side.
+  for (size_t len : {size_t{2}, num}) {
+    for (size_t pos = 0; pos < len; ++pos) {
+      for (size_t i = 0; i < len; ++i) keys[i] = i == pos ? -T{0} : T{0};
+      detail::BaseCase(d, SharedTraits<Traits>(), keys.get(), len, buf.get());
+      size_t count = 0;
+      for (size_t i = 0; i < len; ++i) count += ScalarSignBit(keys[i]);
+      HWY_ASSERT_EQ(size_t{1}, count);
+    }
+  }
+}
+
 HWY_NOINLINE void TestAllBaseCase() {
   // Workaround for stack overflow on MSVC debug.
 #if defined(_MSC_VER)
@@ -317,6 +338,10 @@ HWY_NOINLINE void TestAllBaseCase() {
 
   TestBaseCase<TraitsLane<OrderAscending<int32_t> > >();
   TestBaseCase<TraitsLane<OtherOrder<int64_t> > >();
+  TestBaseCaseSignedZero<TraitsLane<OrderAscending<float>>>();
+#if HWY_HAVE_FLOAT64
+  TestBaseCaseSignedZero<TraitsLane<detail::OrderDescending<double>>>();
+#endif
 #if !HAVE_INTEL
   TestBaseCase<Traits128<OrderAscending128> >();
   TestBaseCase<Traits128<OrderDescending128> >();

hwy/contrib/sort/sorting_networks-inl.h

@@ -51,7 +51,7 @@ struct SharedTraits : public Base {
   HWY_INLINE Vec<D> SortPairsDistance2(D d, Vec<D> v) const {
     const Base* base = static_cast<const Base*>(this);
     Vec<D> swapped = base->SwapAdjacentPairs(d, v);
-    base->Sort2(d, v, swapped);
+    base->template Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenPairs(d, swapped, v);
   }
 
@@ -60,7 +60,7 @@ struct SharedTraits : public Base {
   HWY_INLINE Vec<D> SortPairsReverse8(D d, Vec<D> v) const {
     const Base* base = static_cast<const Base*>(this);
     Vec<D> swapped = base->ReverseKeys8(d, v);
-    base->Sort2(d, v, swapped);
+    base->template Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenQuads(d, swapped, v);
   }
 
@@ -70,7 +70,7 @@ struct SharedTraits : public Base {
     const Base* base = static_cast<const Base*>(this);
     static_assert(Constants::kMaxCols <= 16, "Need actual Reverse16");
     Vec<D> swapped = base->ReverseKeys(d, v);
-    base->Sort2(d, v, swapped);
+    base->template Sort2<PairLayout::kMirrored>(d, v, swapped);
     return ConcatUpperLower(d, swapped, v);  // 8 = half of the vector
   }
 };

hwy/contrib/sort/traits-inl.h

@@ -550,23 +550,32 @@ struct TraitsLane : public Base {
   // according to Base.
   // Corresponds to a conditional swap, which is one "node" of a sorting
   // network. Min/Max are cheaper than compare + blend at least for integers.
-  template <class D>
+  template <PairLayout kLayout = PairLayout::kIndependent, class D>
   HWY_INLINE void Sort2(D d, Vec<D>& a, Vec<D>& b) const {
     const Base* base = static_cast<const Base*>(this);
 
     const Vec<D> a_copy = a;
-    // Prior to AVX3, there is no native 64-bit Min/Max, so they compile to 4
-    // instructions. We can reduce it to a compare + 2 IfThenElse.
+    // Prior to AVX3 there is no native 64-bit integer Min/Max. Use a compare
+    // and two blends, except for mirrored floats: selecting opposite inputs
+    // on a tie there would duplicate one zero after the lanes are merged.
 #if HWY_AVX3 < HWY_TARGET && HWY_TARGET <= HWY_SSSE3
-    if (sizeof(TFromD<D>) == 8) {
+    if (sizeof(TFromD<D>) == 8 &&
+        !(kLayout == PairLayout::kMirrored && IsFloat<TFromD<D>>())) {
       const Mask<D> cmp = base->Compare(d, a, b);
       a = IfThenElse(cmp, a, b);
       b = IfThenElse(cmp, b, a_copy);
       return;
     }
 #endif
     a = base->First(d, a, b);
-    b = base->Last(d, a_copy, b);
+    // Some floating-point Min/Max instructions both choose the second operand
+    // for equal zeros. Keep one of each input. Mirrored lanes already compare
+    // the same pair in opposite orders before selecting one result per lane.
+    if constexpr (kLayout == PairLayout::kIndependent && IsFloat<TFromD<D>>()) {
+      b = base->Last(d, b, a_copy);
+    } else {
+      b = base->Last(d, a_copy, b);
+    }
   }
 
   // Conditionally swaps even-numbered lanes with their odd-numbered neighbor.
@@ -580,7 +589,7 @@ struct TraitsLane : public Base {
     const Vec<D> cmp = VecFromMask(d, base->Compare(d, v, swapped));
     return IfVecThenElse(DupOdd(cmp), swapped, v);
 #else
-    Sort2(d, v, swapped);
+    Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenKeys(swapped, v);
 #endif
   }
@@ -590,7 +599,7 @@ struct TraitsLane : public Base {
   HWY_INLINE Vec<D> SortPairsDistance1(D d, Vec<D> v) const {
     const Base* base = static_cast<const Base*>(this);
     Vec<D> swapped = base->ReverseKeys2(d, v);
-    Sort2(d, v, swapped);
+    Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenKeys(swapped, v);
   }
 
@@ -599,7 +608,7 @@ struct TraitsLane : public Base {
   HWY_INLINE Vec<D> SortPairsReverse4(D d, Vec<D> v) const {
     const Base* base = static_cast<const Base*>(this);
     Vec<D> swapped = base->ReverseKeys4(d, v);
-    Sort2(d, v, swapped);
+    Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenPairs(d, swapped, v);
   }
 
@@ -610,7 +619,7 @@ struct TraitsLane : public Base {
     Vec<D> swapped = base->SwapAdjacentQuads(d, v);
     // Only used in Merge16, so this will not be used on AVX2 (which only has 4
     // u64 lanes), so skip the above optimization for 64-bit AVX2.
-    Sort2(d, v, swapped);
+    Sort2<PairLayout::kMirrored>(d, v, swapped);
     return base->OddEvenQuads(d, swapped, v);
   }
 };

hwy/contrib/sort/traits128-inl.h

@@ -494,7 +494,7 @@ struct Traits128 : public Base {
     return v;
   }
 
-  template <class D, HWY_IF_U64_D(D)>
+  template <PairLayout = PairLayout::kIndependent, class D, HWY_IF_U64_D(D)>
   HWY_INLINE void Sort2(D d, Vec<D>& a, Vec<D>& b) const {
     const Base* base = static_cast<const Base*>(this);