From c03ba3262490d172fedf9f1177f55229c63c31ae Mon Sep 17 00:00:00 2001 From: VirdhatchaniKN Date: Mon, 10 Feb 2025 10:09:27 +0000 Subject: [PATCH] #17758: Switch reader buffer index to compile-time args --- .../device/kernels/dataflow/reader_running_statistics.cpp | 6 +++--- .../device/running_statistics_program_factory.cpp | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp b/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp index e3c457c13c6d..02437e03d6ef 100644 --- a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/kernels/dataflow/reader_running_statistics.cpp +++ b/ttnn/cpp/ttnn/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/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/running_statistics_program_factory.cpp b/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/running_statistics_program_factory.cpp index a8795ae63eb1..a4d6ee3f27c2 100644 --- a/ttnn/cpp/ttnn/operations/normalization/batch_norm/device/running_statistics_program_factory.cpp +++ b/ttnn/cpp/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;