From e9bd503d7d2ea7a2c11836dcf1778990a988a54a Mon Sep 17 00:00:00 2001
From: Allan Liu <aliu@tenstorrent.com>
Date: Mon, 11 Mar 2024 00:46:34 +0000
Subject: [PATCH] #5398: fix unicast binaries

---
 .../watcher/test_waypoint.cpp                 | 44 +++++++------------
 tt_metal/impl/program/program.cpp             |  7 ++-
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/tests/tt_metal/tt_metal/unit_tests_common/watcher/test_waypoint.cpp b/tests/tt_metal/tt_metal/unit_tests_common/watcher/test_waypoint.cpp
index dba1e9f7458..4fcc844e8f2 100644
--- a/tests/tt_metal/tt_metal/unit_tests_common/watcher/test_waypoint.cpp
+++ b/tests/tt_metal/tt_metal/unit_tests_common/watcher/test_waypoint.cpp
@@ -52,26 +52,6 @@ static void RunTest(WatcherFixture* fixture, Device* device) {
         CoreRange(xy_start, xy_end),
         ComputeConfig{});
 
-    // Also run on ethernet cores if they're present
-    bool has_eth_cores = !device->get_active_ethernet_cores(true).empty();
-    // TODO: Change back to a single erisc_kid once the bug with CoreRange fast dispatch is fixed.
-    vector<KernelHandle> erisc_kids;
-    if (has_eth_cores) {
-        std::set<CoreRange> eth_core_ranges;
-        for (const auto& core : device->get_active_ethernet_cores(true)) {
-            eth_core_ranges.insert(CoreRange(core, core));
-            auto erisc_kid = CreateKernel(
-                program,
-                "tests/tt_metal/tt_metal/test_kernels/misc/watcher_waypoints.cpp",
-                core,
-                tt_metal::EthernetConfig{
-                    .noc = tt_metal::NOC::NOC_0
-                }
-            );
-            erisc_kids.push_back(erisc_kid);
-        }
-    }
-
     // The kernels need arguments to be passed in: the number of cycles to delay while syncing,
     // and an L1 buffer to use for the syncing.
     uint32_t clk_mhz = tt::Cluster::instance().get_device_aiclk(device->id());
@@ -108,16 +88,26 @@ static void RunTest(WatcherFixture* fixture, Device* device) {
             );
         }
     }
-    int idx = 0;
-    for (const auto& core : device->get_active_ethernet_cores(true)) {
-        SetRuntimeArgs(
+    // Also run on ethernet cores if they're present
+    bool has_eth_cores = !device->get_active_ethernet_cores(true).empty();
+    if (has_eth_cores) {
+        KernelHandle erisc_kid;
+        std::set<CoreRange> eth_core_ranges;
+        for (const auto& core : device->get_active_ethernet_cores(true)) {
+            eth_core_ranges.insert(CoreRange(core, core));
+        }
+        erisc_kid = CreateKernel(
             program,
-            erisc_kids[idx++],
-            core,
-            args
-        );
+            "tests/tt_metal/tt_metal/test_kernels/misc/watcher_waypoints.cpp",
+            eth_core_ranges,
+            tt_metal::EthernetConfig{.noc = tt_metal::NOC::NOC_0});
+
+        for (const auto& core : device->get_active_ethernet_cores(true)) {
+            SetRuntimeArgs(program, erisc_kid, core, args);
+        }
     }
 
+
     // Run the program in a new thread, we'll have to update gate values in this thread.
     fixture->RunProgram(device, program);
 
diff --git a/tt_metal/impl/program/program.cpp b/tt_metal/impl/program/program.cpp
index 7dd64d11dbf..866d7efd5c0 100644
--- a/tt_metal/impl/program/program.cpp
+++ b/tt_metal/impl/program/program.cpp
@@ -947,7 +947,12 @@ ProgramDeviceMap ConstructProgramDeviceMap(const Device* device, Program& progra
 
     align_program_page_idx_to_new_page();
     for (const KernelGroup& kernel_group : kernel_group_unicast) {
-        populate_program_binaries_pages(kernel_group);
+        if (kernel_group.get_core_type() == CoreType::ETH) {
+            auto kernel = detail::GetKernel(program, kernel_group.erisc_id.value());
+            for (const auto& logical_core : kernel->logical_cores()) {
+                populate_program_binaries_pages(kernel_group);
+            }
+        }
     }
 
     // Since GO signal begin in a new page, I need to advance my idx