From 46f1ed50a270ce3ad51d13b170518f6a6a8221f5 Mon Sep 17 00:00:00 2001 From: VirdhatchaniKN Date: Thu, 27 Feb 2025 13:10:49 +0000 Subject: [PATCH] #18332: Update BN Kernel --- .../batch_norm/device/kernels/compute/batch_norm_kernel.cpp | 5 ++++- .../device/kernels/compute/batch_norm_sfpu_kernel.cpp | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp b/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp index 9ffbdeb1144..c72e0534674 100644 --- a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp +++ b/ttnn/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/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp b/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp index 007ed3e92ae..81323ef7c01 100644 --- a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp +++ b/ttnn/cpp/ttnn/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