#13385: Complete kernel of ttnn.round

tenstorrent · Oct 27, 2024 · 5e5f867 · 5e5f867
1 parent 8a5be87
commit 5e5f867
Showing 1 changed file with 14 additions and 16 deletions.
diff --git a/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_round.h b/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_round.h
@@ -26,24 +26,22 @@ inline vFloat round_even(vFloat v) {
     return v;
 }
 
+static float pow_unsigned(float x, uint n) {
+    if (n == 0)
+        return 1.0f;
+
+    if ((n & 1) == 0)
+        return pow_unsigned(x * x, n >> 1);
+
+    return x * pow_unsigned(x * x, n >> 1);
+}
+
 template <bool APPROXIMATE, int ITERATIONS = 8>
 void calculate_round(const int decimals) {
-    const auto powu = [](vFloat base, uint exp) {
-        vFloat result = 1.0f;
-
-        while (exp > 0) {
-            if (exp & 1)
-                result *= base;
-            base *= base;
-            exp >>= 1;
-        }
-        return result;
-    };
-
-    const auto exp10i = [powu](int exp) {
-        if (exp < 0)
-            return powu(0.1f, -exp);
-        return powu(10.0f, exp);
+    const auto exp10i = [](int n) {
+        if (n < 0)
+            return pow_unsigned(0.1f, -n);
+        return pow_unsigned(10.0f, 0);
     };
 
     const vFloat coeff = exp10i(decimals);