Fix root and cbrt

kfrlib · Oct 11, 2023 · c95fc45 · c95fc45
1 parent 3ec664d
commit c95fc45
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/include/kfr/math/impl/log_exp.hpp b/include/kfr/math/impl/log_exp.hpp
@@ -244,13 +244,19 @@ KFR_INTRINSIC vec<T, N> pow(const vec<T, N>& a, const vec<T, N>& b)
 template <typename T, size_t N>
 KFR_INTRINSIC vec<T, N> root(const vec<T, N>& x, const vec<T, N>& b)
 {
-    return exp(reciprocal(b) * log(x));
+    const vec<T, N> t       = exp(reciprocal(b) * log(abs(x)));
+    const mask<T, N> isint  = floor(b) == b;
+    const mask<T, N> iseven = (broadcastto<itype<T>>(b) & 1) == 0;
+    return select(x > T(), t,
+                  select(x == T(), T(),
+                         select(isint, select(iseven, broadcast<N>(constants<T>::qnan), -t),
+                                broadcast<N>(constants<T>::qnan))));
 }
 
 template <typename T, size_t N>
 KFR_INTRINSIC vec<T, N> cbrt(const vec<T, N>& x)
 {
-    return pow<T, N>(x, T(0.333333333333333333333333333333333));
+    return root<T, N>(x, T(3));
 }
 
 template <typename T, size_t N, KFR_ENABLE_IF(!is_f_class<T>), typename Tout = flt_type<T>>

diff --git a/tests/unit/math/log_exp.cpp b/tests/unit/math/log_exp.cpp
@@ -19,5 +19,15 @@ KFR_AUTO_TEST_1(log, narrow, 2, 1)
 KFR_AUTO_TEST_1(log2, narrow, 2, 1)
 KFR_AUTO_TEST_1(log10, narrow, 3, 1)
 KFR_AUTO_TEST_1(cbrt, narrow, 5, 1)
+
+TEST(cbrt)
+{
+    CHECK(kfr::cbrt(27.f) == 3.f);
+    CHECK(kfr::cbrt(-27.f) == -3.f);
+
+    CHECK(kfr::root(32.f, 5.f) == 2.f);
+    CHECK(kfr::root(-32.f, 5.f) == -2.f);
+    CHECK(std::isnan(kfr::root(-32.f, 5.001f)));
+}
 } // namespace CMT_ARCH_NAME
 } // namespace kfr