Merge pull request google#2458 from johnplatts:hwy_get_exp_enh_013125

PiperOrigin-RevId: 721763477

Author: copybara-github

g3doc/quick_reference.md

@@ -674,6 +674,10 @@ from left to right, of the arguments passed to `Create{2-4}`.
     <code>V **GetExponent**(V v)</code>: returns the exponent of `v[i]` as a floating point value.
     Essentially calculates `floor(log2(x))`.
 
+*   `V`: `{f}`, `VU`: `Vec<RebindToUnsigned<DFromV<V>>>` \
+    <code>VU **GetBiasedExponent**(V v)</code>: returns the biased exponent of
+    `v[i]` as an unsigned integer value.
+
 #### Min/Max
 
 **Note**: Min/Max corner cases are target-specific and may change. If either

hwy/ops/generic_ops-inl.h

@@ -1249,6 +1249,27 @@ HWY_API V MulByFloorPow2(V v, V exp) {
 
 #endif  // HWY_NATIVE_MUL_BY_POW2
 
+// ------------------------------ GetBiasedExponent
+#if (defined(HWY_NATIVE_GET_BIASED_EXPONENT) == defined(HWY_TARGET_TOGGLE))
+#ifdef HWY_NATIVE_GET_BIASED_EXPONENT
+#undef HWY_NATIVE_GET_BIASED_EXPONENT
+#else
+#define HWY_NATIVE_GET_BIASED_EXPONENT
+#endif
+
+template <class V, HWY_IF_FLOAT_V(V)>
+HWY_API VFromD<RebindToUnsigned<DFromV<V>>> GetBiasedExponent(V v) {
+  using T = TFromV<V>;
+
+  const DFromV<V> d;
+  const RebindToUnsigned<decltype(d)> du;
+
+  constexpr int kNumOfMantBits = MantissaBits<T>();
+  return ShiftRight<kNumOfMantBits>(BitCast(du, Abs(v)));
+}
+
+#endif
+
 // ------------------------------ GetExponent
 
 #if (defined(HWY_NATIVE_GET_EXPONENT) == defined(HWY_TARGET_TOGGLE))
@@ -1262,14 +1283,12 @@ template <class V, HWY_IF_FLOAT_V(V)>
 HWY_API V GetExponent(V v) {
   const DFromV<V> d;
   using T = TFromV<V>;
-  const RebindToUnsigned<decltype(d)> du;
   const RebindToSigned<decltype(d)> di;
 
-  constexpr uint8_t mantissa_bits = MantissaBits<T>();
   const auto exponent_offset = Set(di, MaxExponentField<T>() >> 1);
 
   // extract exponent bits as integer
-  const auto encoded_exponent = ShiftRight<mantissa_bits>(BitCast(du, Abs(v)));
+  const auto encoded_exponent = GetBiasedExponent(v);
   const auto exponent_int = Sub(BitCast(di, encoded_exponent), exponent_offset);
 
   // convert integer to original type

hwy/ops/ppc_vsx-inl.h

@@ -1939,9 +1939,6 @@ HWY_API Vec128<T, N> ApproximateReciprocal(Vec128<T, N> v) {
 #endif
 }
 
-// TODO: Implement GetExponent using vec_extract_exp (which returns the biased
-//       exponent) followed by a subtraction by MaxExponentField<T>() >> 1
-
 // ------------------------------ Floating-point square root
 
 #if HWY_S390X_HAVE_Z14
@@ -1979,6 +1976,23 @@ HWY_API Vec128<T, N> Sqrt(Vec128<T, N> v) {
   return Vec128<T, N>{vec_sqrt(v.raw)};
 }
 
+// ------------------------------ GetBiasedExponent
+
+#if HWY_PPC_HAVE_9
+
+#ifdef HWY_NATIVE_GET_BIASED_EXPONENT
+#undef HWY_NATIVE_GET_BIASED_EXPONENT
+#else
+#define HWY_NATIVE_GET_BIASED_EXPONENT
+#endif
+
+template <class V, HWY_IF_FLOAT3264_V(V)>
+HWY_API VFromD<RebindToUnsigned<DFromV<V>>> GetBiasedExponent(V v) {
+  return VFromD<RebindToUnsigned<DFromV<V>>>{vec_extract_exp(v.raw)};
+}
+
+#endif  // HWY_PPC_HAVE_9
+
 // ------------------------------ Min (Gt, IfThenElse)
 
 template <typename T, size_t N, HWY_IF_NOT_SPECIAL_FLOAT(T)>

hwy/ops/shared-inl.h

@@ -631,8 +631,11 @@ HWY_API bool IsAligned(D d, T* ptr) {
 #define HWY_IF_SIGNED_V(V) HWY_IF_SIGNED(hwy::HWY_NAMESPACE::TFromV<V>)
 #define HWY_IF_FLOAT_V(V) HWY_IF_FLOAT(hwy::HWY_NAMESPACE::TFromV<V>)
 #define HWY_IF_NOT_FLOAT_V(V) HWY_IF_NOT_FLOAT(hwy::HWY_NAMESPACE::TFromV<V>)
+#define HWY_IF_FLOAT3264_V(V) HWY_IF_FLOAT3264(hwy::HWY_NAMESPACE::TFromV<V>)
 #define HWY_IF_SPECIAL_FLOAT_V(V) \
   HWY_IF_SPECIAL_FLOAT(hwy::HWY_NAMESPACE::TFromV<V>)
+#define HWY_IF_FLOAT_OR_SPECIAL_V(V) \
+  HWY_IF_FLOAT_OR_SPECIAL(hwy::HWY_NAMESPACE::TFromV<V>)
 #define HWY_IF_NOT_FLOAT_NOR_SPECIAL_V(V) \
   HWY_IF_NOT_FLOAT_NOR_SPECIAL(hwy::HWY_NAMESPACE::TFromV<V>)
 

hwy/ops/x86_128-inl.h

@@ -5170,6 +5170,33 @@ HWY_API V AbsDiff(V a, V b) {
   return Abs(a - b);
 }
 
+// ------------------------------ GetExponent
+
+#if HWY_TARGET <= HWY_AVX3
+
+#ifdef HWY_NATIVE_GET_EXPONENT
+#undef HWY_NATIVE_GET_EXPONENT
+#else
+#define HWY_NATIVE_GET_EXPONENT
+#endif
+
+#if HWY_HAVE_FLOAT16
+template <class V, HWY_IF_F16(TFromV<V>), HWY_IF_V_SIZE_LE_V(V, 16)>
+HWY_API V GetExponent(V v) {
+  return V{_mm_getexp_ph(v.raw)};
+}
+#endif
+template <class V, HWY_IF_F32(TFromV<V>), HWY_IF_V_SIZE_LE_V(V, 16)>
+HWY_API V GetExponent(V v) {
+  return V{_mm_getexp_ps(v.raw)};
+}
+template <class V, HWY_IF_F64(TFromV<V>), HWY_IF_V_SIZE_LE_V(V, 16)>
+HWY_API V GetExponent(V v) {
+  return V{_mm_getexp_pd(v.raw)};
+}
+
+#endif
+
 // ------------------------------ MaskedMinOr
 
 #if HWY_TARGET <= HWY_AVX3

hwy/ops/x86_256-inl.h

@@ -2734,6 +2734,27 @@ HWY_API Vec256<double> ApproximateReciprocal(Vec256<double> v) {
 }
 #endif
 
+// ------------------------------ GetExponent
+
+#if HWY_TARGET <= HWY_AVX3
+
+#if HWY_HAVE_FLOAT16
+template <class V, HWY_IF_F16(TFromV<V>), HWY_IF_V_SIZE_V(V, 32)>
+HWY_API V GetExponent(V v) {
+  return V{_mm256_getexp_ph(v.raw)};
+}
+#endif
+template <class V, HWY_IF_F32(TFromV<V>), HWY_IF_V_SIZE_V(V, 32)>
+HWY_API V GetExponent(V v) {
+  return V{_mm256_getexp_ps(v.raw)};
+}
+template <class V, HWY_IF_F64(TFromV<V>), HWY_IF_V_SIZE_V(V, 32)>
+HWY_API V GetExponent(V v) {
+  return V{_mm256_getexp_pd(v.raw)};
+}
+
+#endif
+
 // ------------------------------ MaskedMinOr
 
 #if HWY_TARGET <= HWY_AVX3

hwy/ops/x86_512-inl.h

@@ -1842,7 +1842,22 @@ HWY_API Vec512<double> ApproximateReciprocal(Vec512<double> v) {
   return Vec512<double>{_mm512_rcp14_pd(v.raw)};
 }
 
-// TODO: Implement GetExponent using _mm_getexp_ps/_mm_getexp_pd/_mm_getexp_ph
+// ------------------------------ GetExponent
+
+#if HWY_HAVE_FLOAT16
+template <class V, HWY_IF_F16(TFromV<V>), HWY_IF_V_SIZE_V(V, 64)>
+HWY_API V GetExponent(V v) {
+  return V{_mm512_getexp_ph(v.raw)};
+}
+#endif
+template <class V, HWY_IF_F32(TFromV<V>), HWY_IF_V_SIZE_V(V, 64)>
+HWY_API V GetExponent(V v) {
+  return V{_mm512_getexp_ps(v.raw)};
+}
+template <class V, HWY_IF_F64(TFromV<V>), HWY_IF_V_SIZE_V(V, 64)>
+HWY_API V GetExponent(V v) {
+  return V{_mm512_getexp_pd(v.raw)};
+}
 
 // ------------------------------ MaskedMinOr
 

hwy/tests/float_test.cc

@@ -605,18 +605,30 @@ HWY_NOINLINE void TestAllAbsDiff() {
 struct TestGetExponent {
   template <typename T, class D>
   HWY_NOINLINE void operator()(T /*unused*/, D d) {
+    const RebindToUnsigned<decltype(d)> du;
+
+    using TFArith = If<IsSpecialFloat<T>(), float, T>;
+    using TU = MakeUnsigned<T>;
     const size_t N = Lanes(d);
 
     auto v = Iota(d, 1);
 
     auto expected = AllocateAligned<T>(N);
-    HWY_ASSERT(expected);
+    auto expected_biased = AllocateAligned<TU>(N);
+    HWY_ASSERT(expected && expected_biased);
+
+    constexpr int kNumOfMantBits = MantissaBits<T>();
 
     for (size_t i = 0; i < N; ++i) {
-      auto test_val = (float)(i + 1);
-      expected[i] = ConvertScalarTo<T>(std::floor(std::log2(test_val)));
+      const T test_val = ConvertScalarTo<T>(i + 1);
+      expected[i] = ConvertScalarTo<T>(
+          std::floor(std::log2(ConvertScalarTo<TFArith>(test_val))));
+      expected_biased[i] =
+          static_cast<TU>((BitCastScalar<TU>(test_val) >> kNumOfMantBits) &
+                          static_cast<TU>(MaxExponentField<T>()));
     }
     HWY_ASSERT_VEC_EQ(d, expected.get(), GetExponent(v));
+    HWY_ASSERT_VEC_EQ(du, expected_biased.get(), GetBiasedExponent(v));
   }
 };