fix: Remove epsilon from the comparisons in NE/CPPTopKV

The comparison in CPP and NETopKV functions don't use type-dependent epsilon anymore. A simple comparison.

This is categorized as fix because it's a very minor behavioural change in some edge-case numerical situations.

Partially Resolves: MLCE-1719

Change-Id: Ic0d02f6df05ce45886bcee458c2bcd578a388ee2
Signed-off-by: Gunes Bayir <gunes.bayir@arm.com>

Author: gunes-arm

src/core/CPP/kernels/CPPTopKVKernel.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020 Arm Limited.
+ * Copyright (c) 2019-2020, 2026 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -34,14 +34,7 @@ namespace arm_compute
 {
 namespace
 {
-template <typename T, typename std::enable_if<utils::traits::is_floating_point<T>::value, int>::type = 0>
-inline bool greater_than(T a, T b)
-{
-    const T epsilon = std::numeric_limits<T>::epsilon();
-    return (a - b > epsilon);
-}
-
-template <typename T, typename std::enable_if<!utils::traits::is_floating_point<T>::value, int>::type = 0>
+template <typename T>
 inline bool greater_than(T a, T b)
 {
     return (a > b);

src/core/NEON/wrapper/intrinsics/shr.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Arm Limited.
+ * Copyright (c) 2022, 2026 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -22,8 +22,8 @@
  * SOFTWARE.
  */
 
-#ifndef ARM_COMPUTE_WRAPPER_SHR_H
-#define ARM_COMPUTE_WRAPPER_SHR_H
+#ifndef ACL_SRC_CORE_NEON_WRAPPER_INTRINSICS_SHR_H
+#define ACL_SRC_CORE_NEON_WRAPPER_INTRINSICS_SHR_H
 
 #include <arm_neon.h>
 #include <type_traits>
@@ -103,6 +103,8 @@ VSHR_IMPL(int8x8_t, vshr_n, s8)
     {                                      \
         return prefix##_##postfix(a, b);   \
     }
+VSHRQ_IMPL(uint16x8_t, vshrq_n, u16)
+VSHRQ_IMPL(int16x8_t, vshrq_n, s16)
 VSHRQ_IMPL(uint32x4_t, vshrq_n, u32)
 VSHRQ_IMPL(int32x4_t, vshrq_n, s32)
 #undef VSHRQ_IMPL
@@ -145,4 +147,4 @@ VQRSHRN_EX_SCALAR_IMPL(int32_t, int64_t, vqrshrnd_n, vqrshrund_n, s64)
 
 } // namespace wrapper
 } // namespace arm_compute
-#endif /* ARM_COMPUTE_WRAPPER_SHR_H */
+#endif // ACL_SRC_CORE_NEON_WRAPPER_INTRINSICS_SHR_H

src/cpu/kernels/topkv/generic/neon/fp16.cpp

@@ -23,6 +23,7 @@
  */
 #if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && defined(ENABLE_FP16_KERNELS)
 
+#include "src/core/NEON/wrapper/wrapper.h"
 #include "src/cpu/kernels/topkv/generic/neon/impl.h"
 
 #include <arm_neon.h>
@@ -33,33 +34,24 @@ namespace cpu
 {
 namespace detail
 {
-static inline uint32_t reduce_u32x4(uint32x4_t v)
+static inline uint32_t reduce_u16x8(uint16x8_t v)
 {
 #if defined(__aarch64__)
-    return vaddvq_u32(v);
+    return vaddvq_u16(v);
 #else
-    uint32x2_t s = vadd_u32(vget_low_u32(v), vget_high_u32(v));
-    s            = vpadd_u32(s, s);
-    return vget_lane_u32(s, 0);
+    uint16x4_t s = vadd_u16(vget_low_u16(v), vget_high_u16(v));
+    s            = vpadd_u16(s, s);
+    return vget_lane_u16(s, 0);
 #endif
 }
 
-// Explicit specialization for float16_t
 template <>
-uint32_t count_gt_block<float16_t>(const float16_t *ptr, float16_t threshold)
+uint32_t count_gt_block<float16_t>(const float16_t *ptr, float16x8_t thr_vec)
 {
-    // Load 8 fp16
-    const float16x8_t v16 = vld1q_f16(reinterpret_cast<const __fp16 *>(ptr));
-
-    const float32x4_t thr = vdupq_n_f32(threshold);
-
-    const float32x4_t v0 = vcvt_f32_f16(vget_low_f16(v16));
-    const float32x4_t v1 = vcvt_f32_f16(vget_high_f16(v16));
-
-    const uint32x4_t b0 = vshrq_n_u32(vcgtq_f32(v0, thr), 31);
-    const uint32x4_t b1 = vshrq_n_u32(vcgtq_f32(v1, thr), 31);
-
-    return reduce_u32x4(b0) + reduce_u32x4(b1);
+    const float16x8_t v    = wrapper::vloadq(ptr);
+    const uint16x8_t  mask = wrapper::vcgt(v, thr_vec); // new: v > (threshold)
+    const uint16x8_t  b    = wrapper::vshrq_n<15>(mask);
+    return reduce_u16x8(b);
 }
 
 } // namespace detail

src/cpu/kernels/topkv/generic/neon/fp32.cpp

@@ -47,20 +47,11 @@ static inline uint32_t reduce_u32x4(uint32x4_t v)
 
 // Explicit specialization for float: may use float32x4_t
 template <>
-uint32_t count_gt_block<float>(const float *ptr, float threshold)
+uint32_t count_gt_block<float>(const float *ptr, float32x4_t thr_vec)
 {
-    using Tag = wrapper::traits::neon_bitvector_tag_t<float, wrapper::traits::BitWidth::W128>;
-
-    const auto v = wrapper::vloadq(ptr);
-
-    // epsilon-aware compare: treat a > b only when (a - b) > epsilon
-    const float eps_val      = std::numeric_limits<float>::epsilon();
-    const float thr_with_eps = threshold + eps_val;
-    const auto  thr_eps_vec  = wrapper::vdup_n(thr_with_eps, Tag{});
-    const auto  mask         = wrapper::vcgt(v, thr_eps_vec); // new: v > (threshold + eps)
-
-    const uint32x4_t m = mask;
-    const uint32x4_t b = vshrq_n_u32(m, 31);
+    const float32x4_t v    = wrapper::vloadq(ptr);
+    const uint32x4_t  mask = wrapper::vcgt(v, thr_vec); // new: v > (threshold)
+    const uint32x4_t  b    = wrapper::vshrq_n<31>(mask);
     return reduce_u32x4(b);
 }
 } // namespace detail

src/cpu/kernels/topkv/generic/neon/impl.h

@@ -27,6 +27,8 @@
 #include "arm_compute/core/Helpers.h"
 #include "arm_compute/core/Window.h"
 
+#include "src/core/NEON/wrapper/wrapper.h"
+
 #include <limits>
 #include <type_traits>
 
@@ -36,13 +38,18 @@ namespace cpu
 {
 namespace detail
 {
-template <typename ScalarType>
-uint32_t count_gt_block(const ScalarType *ptr, ScalarType threshold);
+template <typename ScalarType, typename VectorType>
+uint32_t count_gt_block(const ScalarType *ptr, VectorType threshold);
 
 template <typename ScalarType>
 void topkv_neon_wrapper(
     const ITensor *predictions, const ITensor *targets, ITensor *out, uint32_t k, const Window &window)
 {
+    constexpr auto bit_width = wrapper::traits::BitWidth::W128;
+
+    using TagType    = typename wrapper::traits::neon_bitvector_tag_t<ScalarType, bit_width>;
+    using VectorType = typename wrapper::traits::neon_bitvector_t<ScalarType, bit_width>;
+
     const auto        &pred_info = *predictions->info();
     const unsigned int C         = pred_info.tensor_shape()[0];
 
@@ -66,7 +73,8 @@ void topkv_neon_wrapper(
             const ScalarType *base =
                 reinterpret_cast<const ScalarType *>(predictions->ptr_to_element(Coordinates{0, n}));
 
-            const ScalarType thr = base[t];
+            const ScalarType thr     = base[t];
+            const VectorType thr_vec = wrapper::vdup_n(thr, TagType{});
 
             uint32_t     rank = 0;
             unsigned int c    = 0;
@@ -75,7 +83,7 @@ void topkv_neon_wrapper(
             // Vector loop with early-exit
             for (; c + vec_elems <= C; c += vec_elems)
             {
-                rank += count_gt_block<ScalarType>(base + c, thr);
+                rank += count_gt_block<ScalarType>(base + c, thr_vec);
                 if (rank >= k)
                 {
                     // For large C and small K (e.g. QASYMM8, C=32000, K=3), the probability that the

src/cpu/kernels/topkv/generic/neon/integer.cpp

@@ -43,12 +43,11 @@ static inline uint32_t reduce_u32x4(uint32x4_t v)
 }
 
 template <>
-uint32_t count_gt_block<int32_t>(const int32_t *ptr, int32_t threshold)
+uint32_t count_gt_block<int32_t>(const int32_t *ptr, int32x4_t thr_vec)
 {
-    const int32x4_t  v   = vld1q_s32(ptr);
-    const int32x4_t  thr = vdupq_n_s32(threshold);
-    const uint32x4_t m   = vcgtq_s32(v, thr);  // 0xFFFFFFFF / 0 per lane
-    const uint32x4_t b   = vshrq_n_u32(m, 31); // 0/1 per lane
+    const int32x4_t  v = vld1q_s32(ptr);
+    const uint32x4_t m = vcgtq_s32(v, thr_vec); // 0xFFFFFFFF / 0 per lane
+    const uint32x4_t b = vshrq_n_u32(m, 31);    // 0/1 per lane
     return reduce_u32x4(b);
 }
 

src/cpu/kernels/topkv/generic/neon/qasymm8.cpp

@@ -48,11 +48,10 @@ static inline uint32_t reduce_u8_to_count(uint8x16_t m)
 }
 
 template <>
-uint32_t count_gt_block<uint8_t>(const uint8_t *ptr, uint8_t threshold)
+uint32_t count_gt_block<uint8_t>(const uint8_t *ptr, uint8x16_t thr_vec)
 {
-    const uint8x16_t v   = vld1q_u8(ptr);
-    const uint8x16_t thr = vdupq_n_u8(threshold);
-    const uint8x16_t m   = vcgtq_u8(v, thr); // 0xFF / 0x00 bytes
+    const uint8x16_t v = vld1q_u8(ptr);
+    const uint8x16_t m = vcgtq_u8(v, thr_vec); // 0xFF / 0x00 bytes
     return reduce_u8_to_count(m);
 }
 

src/cpu/kernels/topkv/generic/neon/qasymm8_signed.cpp

@@ -46,11 +46,10 @@ static inline uint32_t reduce_u8_to_count(uint8x16_t m)
 }
 
 template <>
-uint32_t count_gt_block<int8_t>(const int8_t *ptr, int8_t threshold)
+uint32_t count_gt_block<int8_t>(const int8_t *ptr, int8x16_t thr_vec)
 {
-    const int8x16_t  v   = vld1q_s8(ptr);
-    const int8x16_t  thr = vdupq_n_s8(threshold);
-    const uint8x16_t m   = vcgtq_s8(v, thr); // returns uint8x16_t mask
+    const int8x16_t  v = vld1q_s8(ptr);
+    const uint8x16_t m = vcgtq_s8(v, thr_vec); // returns uint8x16_t mask
     return reduce_u8_to_count(m);
 }
 

tests/validation/reference/TopKV.cpp

@@ -50,24 +50,21 @@ SimpleTensor<uint8_t> topkv(SimpleTensor<T> &predictions, SimpleTensor<uint32_t>
 
     SimpleTensor<uint8_t> expected(TensorShape(N), DataType::U8);
 
-    const float eps = std::numeric_limits<float>::epsilon();
-
     for (int i = 0; i < N; ++i)
     {
         // targets[i] (U32)
         const uint32_t target_class = targets[i];
 
         // Read predictions[target_class, i] as T, then promote to float
-        const T     target_t   = *reinterpret_cast<const T *>(predictions(Coordinates{target_class, i}));
-        const float target_val = static_cast<float>(target_t);
+        const T target    = *reinterpret_cast<const T *>(predictions(Coordinates{target_class, i}));
+        const T threshold = target;
 
         unsigned int rank = 0;
         for (int c = 0; c < C; ++c)
         {
-            const T     vt = *reinterpret_cast<const T *>(predictions(Coordinates{c, i}));
-            const float v  = static_cast<float>(vt);
+            const T v = *reinterpret_cast<const T *>(predictions(Coordinates{c, i}));
 
-            if ((v - target_val) > eps)
+            if (v > threshold)
             {
                 ++rank;
             }