#18332: Update BN Kernel

tenstorrent · Feb 27, 2025 · 46f1ed5 · 46f1ed5
1 parent be1793d
commit 46f1ed5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/...cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp b/...cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp
@@ -163,7 +163,8 @@ void MAIN {
     uint32_t complete_iterations = (num_tiles + tile_start) / tile_freq;
     uint32_t remaining_iterations = (num_tiles + tile_start) % tile_freq;
 
-    cb_wait_front(cb_eps, 1);
+    constexpr uint32_t onetile = 1;
+    cb_wait_front(cb_eps, onetile);
 
     for (uint32_t i = 0; i < complete_iterations; ++i, tile_start = 0) {
         batchnorm_bcast_tiles(
@@ -199,5 +200,7 @@ void MAIN {
             weight_has_value,
             bias_has_value);
     }
+
+    cb_pop_front(cb_eps, onetile);
 }
 }  // namespace NAMESPACE
diff --git a/...tnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp b/...tnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp
@@ -238,5 +238,7 @@ void MAIN {
             weight_has_value,
             bias_has_value);
     }
+
+    cb_pop_front(cb_eps, onetile);
 }
 }  // namespace NAMESPACE