From 08c27760ff8d88c6514d74d9c59cdb7f32ba441d Mon Sep 17 00:00:00 2001
From: Reem Tawfik <rtawfik@tenstorrent.com>
Date: Wed, 5 Jun 2024 03:04:26 +0000
Subject: [PATCH] #9036: Standardize llk sfpu inits

---
 .../llk_math_eltwise_unary_sfpu_dropout.h     |  2 +-
 .../llk_math_eltwise_unary_sfpu_init.h        | 19 +++++-------------
 .../llk_math_eltwise_unary_sfpu_dropout.h     |  2 +-
 .../llk_math_eltwise_unary_sfpu_init.h        | 20 +++++--------------
 4 files changed, 12 insertions(+), 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
index 4dfddab02ea..23f036d522f 100644
--- 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
index 4565c88949b..2fa29cd57be 100644
--- 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/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h b/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h
index 4dfddab02ea..23f036d522f 100644
--- a/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/llk_math_eltwise_unary_sfpu_dropout.h
+++ b/tt_metal/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
index 4565c88949b..96670f2a45d 100644
--- 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);
-}
-
 }