#9036: Standardize llk sfpu inits

tenstorrent · Jun 6, 2024 · 56e55d5 · 56e55d5
1 parent 2e37c02
commit 56e55d5
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 31 deletions.
diff --git a/tt_metal/hw/ckernels/blackhole/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h b/tt_metal/hw/ckernels/blackhole/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h
@@ -14,7 +14,7 @@ namespace ckernel {
 
 template <bool APPROXIMATE>
 inline void llk_math_eltwise_unary_sfpu_dropout_init(uint seed = 0) {
-    llk_math_eltwise_unary_sfpu_init_1_param<SfpuType::dropout, APPROXIMATE>(sfpu::dropout_init<APPROXIMATE>, seed);
+    llk_math_eltwise_unary_sfpu_init<SfpuType::dropout, APPROXIMATE>(sfpu::dropout_init<APPROXIMATE>, seed);
 }
 
 template <bool APPROXIMATE>

diff --git a/tt_metal/hw/ckernels/blackhole/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_init.h b/tt_metal/hw/ckernels/blackhole/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_init.h
@@ -11,22 +11,13 @@ namespace ckernel {
 
 template <SfpuType sfpu_op, bool APPROXIMATE>
 inline void llk_math_eltwise_unary_sfpu_init() {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    math::reset_counters(p_setrwc::SET_ABD_F);
+    _llk_math_eltwise_unary_sfpu_init_<sfpu_op>();
 }
 
-template <SfpuType sfpu_op, bool APPROXIMATE>
-inline void llk_math_eltwise_unary_sfpu_init(void (*func)()) {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    func();
-    math::reset_counters(p_setrwc::SET_ABD_F);
-}
-
-template <SfpuType sfpu_op, bool APPROXIMATE>
-inline void llk_math_eltwise_unary_sfpu_init_1_param(void (*func)(uint), uint param0 = 0) {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    func(param0);
-    math::reset_counters(p_setrwc::SET_ABD_F);
+template <SfpuType sfpu_op, bool APPROXIMATE, class F, class ... ARGS>
+inline void llk_math_eltwise_unary_sfpu_init(F&& init_func, ARGS&& ... args) {
+    _llk_math_eltwise_unary_sfpu_init_<sfpu_op>();
+    init_func(static_cast<ARGS&&>(args)...);
 }
 
 }
diff --git a/...etal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h b/...etal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h
@@ -14,7 +14,7 @@ namespace ckernel {
 
 template <bool APPROXIMATE>
 inline void llk_math_eltwise_unary_sfpu_dropout_init(uint seed = 0) {
-    llk_math_eltwise_unary_sfpu_init_1_param<SfpuType::dropout, APPROXIMATE>(sfpu::dropout_init<APPROXIMATE>, seed);
+    llk_math_eltwise_unary_sfpu_init<SfpuType::dropout, APPROXIMATE>(sfpu::dropout_init<APPROXIMATE>, seed);
 }
 
 template <bool APPROXIMATE>

diff --git a/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_init.h b/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_init.h
@@ -11,22 +11,12 @@ namespace ckernel {
 
 template <SfpuType sfpu_op, bool APPROXIMATE>
 inline void llk_math_eltwise_unary_sfpu_init() {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    math::reset_counters(p_setrwc::SET_ABD_F);
+    _llk_math_eltwise_unary_sfpu_init_<sfpu_op>();
 }
 
-template <SfpuType sfpu_op, bool APPROXIMATE>
-inline void llk_math_eltwise_unary_sfpu_init(void (*func)()) {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    func();
-    math::reset_counters(p_setrwc::SET_ABD_F);
+template <SfpuType sfpu_op, bool APPROXIMATE, class F, class ... ARGS>
+inline void llk_math_eltwise_unary_sfpu_init(F&& init_func, ARGS&& ... args) {
+    _llk_math_eltwise_unary_sfpu_init_<sfpu_op>();
+    init_func(static_cast<ARGS&&>(args)...);
 }
-
-template <SfpuType sfpu_op, bool APPROXIMATE>
-inline void llk_math_eltwise_unary_sfpu_init_1_param(void (*func)(uint), uint param0 = 0) {
-    eltwise_unary_sfpu_configure_addrmod<sfpu_op>();
-    func(param0);
-    math::reset_counters(p_setrwc::SET_ABD_F);
-}
-
 }