#0: Remove pop for scalar in BN

tests/ttnn/unit_tests/operations/test_batch_norm.py

@@ -19,9 +19,7 @@
     [
         *(torch.Size([n, c, 32, 32]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
         *(torch.Size([n, c, 23, 23]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
-        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2], [1, 2, 3])),
-        torch.Size([3, 1, 64, 120]),
-        torch.Size([3, 2, 64, 120]),
+        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2, 3], [1, 2, 3, 4])),
     ],
 )
 @pytest.mark.parametrize(
@@ -171,9 +169,7 @@ def test_BN_fp32_full_value(device, channel_size, eps, weight, bias):
     [
         *(torch.Size([n, c, 32, 32]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
         *(torch.Size([n, c, 23, 23]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
-        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2], [1, 2, 3])),
-        torch.Size([3, 1, 64, 120]),
-        torch.Size([3, 2, 64, 120]),
+        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2, 3], [1, 2, 3, 4])),
     ],
 )
 @pytest.mark.parametrize(
@@ -248,9 +244,7 @@ def test_batch_norm_fp32(
     [
         *(torch.Size([n, c, 32, 32]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
         *(torch.Size([n, c, 23, 23]) for n, c in product([1, 2, 3, 4], [1, 2, 3, 4])),
-        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2], [1, 2, 3])),
-        torch.Size([3, 1, 64, 120]),
-        torch.Size([3, 2, 64, 120]),
+        *(torch.Size([n, c, 64, 120]) for n, c in product([1, 2, 3], [1, 2, 3, 4])),
     ],
 )
 @pytest.mark.parametrize(

ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_kernel.cpp

@@ -53,7 +53,6 @@ ALWI void batchnorm_bcast_tiles(
     // 1/(sqrt(batch_var + eps))
     cb_reserve_back(cb_den, onetile);
     cb_wait_front(cb_batch_var, 1);
-    cb_wait_front(cb_eps, 1);
 
     tile_regs_acquire();
     add_tiles_init_with_dt(cb_batch_var, cb_eps);
@@ -67,7 +66,6 @@ ALWI void batchnorm_bcast_tiles(
     tile_regs_release();
 
     cb_pop_front(cb_batch_var, 1);
-    cb_pop_front(cb_eps, 1);
     cb_push_back(cb_den, onetile);
 
     // (input - batch_mean)/(sqrt(batch_var + eps)) = result
@@ -164,6 +162,9 @@ void MAIN {
     sub_tiles_init(cb_other, cb_bcast);
     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);
+
     for (uint32_t i = 0; i < complete_iterations; ++i, tile_start = 0) {
         batchnorm_bcast_tiles(
             cb_bcast,
@@ -198,8 +199,5 @@ void MAIN {
             weight_has_value,
             bias_has_value);
     }
-
-    constexpr uint32_t onetile = 1;
-    constexpr int dst0 = 0;
 }
 }  // namespace NAMESPACE

ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/compute/batch_norm_sfpu_kernel.cpp

@@ -63,7 +63,6 @@ ALWI void batchnorm_bcast_tiles(
     // 1/(sqrt(batch_var + eps))
     cb_reserve_back(cb_den, onetile);
     cb_wait_front(cb_batch_var, onetile);
-    cb_wait_front(cb_eps, onetile);
 
     add_binary_tile_init();
     rsqrt_tile_init();
@@ -86,7 +85,6 @@ ALWI void batchnorm_bcast_tiles(
 
     cb_push_back(cb_den, onetile);
     cb_pop_front(cb_batch_var, onetile);
-    cb_pop_front(cb_eps, onetile);
 
     // (input - batch_mean)/(sqrt(batch_var + eps)) = result
     cb_wait_front(cb_den, onetile);
@@ -202,6 +200,10 @@ void MAIN {
 
     uint32_t complete_iterations = (num_tiles + tile_start) / tile_freq;
     uint32_t remaining_iterations = (num_tiles + tile_start) % tile_freq;
+
+    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(
             cb_bcast,
@@ -236,8 +238,5 @@ void MAIN {
             weight_has_value,
             bias_has_value);
     }
-
-    constexpr uint32_t onetile = 1;
-    constexpr int dst0 = 0;
 }
 }  // namespace NAMESPACE