#17758: Switch reader buffer index to compile-time args

tenstorrent · Feb 11, 2025 · c03ba32 · c03ba32
1 parent 30b0aa9
commit c03ba32
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/...operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp b/...operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp
@@ -21,9 +21,9 @@ void kernel_main() {
 
     constexpr bool src_is_dram = get_compile_time_arg_val(0) == 1;
 
-    constexpr auto cb_id_src = tt::CBIndex::c_0;
-    constexpr auto cb_id_momentum = tt::CBIndex::c_5;
-    constexpr auto cb_id_one = tt::CBIndex::c_6;
+    constexpr auto cb_id_src = get_compile_time_arg_val(1);
+    constexpr auto cb_id_momentum = get_compile_time_arg_val(2);
+    constexpr auto cb_id_one = get_compile_time_arg_val(3);
     constexpr uint32_t onetile = 1;
 
     const uint32_t src_tile_bytes = get_tile_size(cb_id_src);

diff --git a/...pp/ttnn/operations/normalization/batch_norm/device/running_statistics_program_factory.cpp b/...pp/ttnn/operations/normalization/batch_norm/device/running_statistics_program_factory.cpp
@@ -261,7 +261,8 @@ RunningStatistics::RunningStatisticsProgramFactory::create(
         program,
         "ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp",
         all_device_cores,
-        tt_metal::ReaderDataMovementConfig({a_is_dram}, std::move(reader_defines)));
+        tt_metal::ReaderDataMovementConfig(
+            {a_is_dram, batch_mean_tensor_cb, momentum_cb, one_cb}, std::move(reader_defines)));
 
     // WRITER KERNEL
     auto writer_defines = dataflow_defines;