#4931: add apis to get ethernet by socket ids

tests/tt_metal/tt_metal/unit_tests/multichip/device_cluster_api.cpp

@@ -36,10 +36,6 @@ TEST_F(N300DeviceFixture, ValidateEthernetConnectivity) {
     ASSERT_TRUE(device_0->get_inactive_ethernet_cores().size() == 14);
     ASSERT_TRUE(device_1->get_inactive_ethernet_cores().size() == 14);
 
-    // Check connectivity between chips
-    ASSERT_TRUE(device_0->get_ethernet_connected_chip_ids() == std::unordered_set({1}));
-    ASSERT_TRUE(device_1->get_ethernet_connected_chip_ids() == std::unordered_set({0}));
-
     for (const auto& core : device_0_active_eth_cores) {
         std::tuple<chip_id_t, CoreCoord> core_on_chip_1 = device_0->get_connected_ethernet_core(core);
         ASSERT_TRUE(std::get<0>(core_on_chip_1) == 1);
@@ -142,4 +138,20 @@ TEST_F(N300DeviceFixture, ValidatePhysicalCoreConversion) {
     // Check an invalid core type
     EXPECT_ANY_THROW(device_0->physical_core_from_logical_core({.x = 0, .y = 0}, CoreType::DRAM));
 }
+
+TEST_F(N300DeviceFixture, ValidateEthernetSockets) {
+    const auto& device_0 = this->devices_.at(0);
+    const auto& device_1 = this->devices_.at(1);
+
+    std::vector<CoreCoord> device_0_sockets = device_0->get_ethernet_sockets(1);
+    std::vector<CoreCoord> device_1_sockets = device_1->get_ethernet_sockets(0);
+
+    ASSERT_TRUE(device_0_sockets.size() == 2);
+    ASSERT_TRUE(device_1_sockets.size() == 2);
+    ASSERT_TRUE(
+        device_0->get_connected_ethernet_core(device_0_sockets.at(0)) == std::make_tuple(device_1->id(), device_1_sockets.at(0)));
+    ASSERT_TRUE(
+        device_0->get_connected_ethernet_core(device_0_sockets.at(1)) == std::make_tuple(device_1->id(), device_1_sockets.at(1)));
+    EXPECT_ANY_THROW(device_0->get_ethernet_sockets(2));
+}
 }  // namespace unit_tests::multichip::cluster

tests/tt_metal/tt_metal/unit_tests/multichip/ring_gather_kernels.cpp

@@ -80,7 +80,7 @@ std::vector<Device*> get_device_ring(std::vector<tt::tt_metal::Device*> devices)
     std::vector<std::vector<int>> adj(devices.size(), std::vector<int>(devices.size(), 0));
     for (uint32_t i = 0; i < devices.size(); ++i) {
         const auto& device = devices[i];
-        for (const auto& connected_device_id : device->get_ethernet_connected_chip_ids()) {
+        for (const auto& [connected_device_id, cores] : device->get_ethernet_cores_grouped_by_connected_chips()) {
             for (uint32_t j = 0; j < devices.size(); ++j) {
                 if (devices[j]->id() == connected_device_id) {
                     adj[i][j] = 1;

tests/tt_metal/tt_metal/unit_tests_fast_dispatch/multichip/test_eth_ring_gather_EnqueueProgram.cpp

@@ -81,7 +81,7 @@ std::vector<Device*> get_device_ring(std::vector<tt::tt_metal::Device*> devices)
     std::vector<std::vector<int>> adj(devices.size(), std::vector<int>(devices.size(), 0));
     for (uint32_t i = 0; i < devices.size(); ++i) {
         const auto& device = devices[i];
-        for (const auto& connected_device_id : device->get_ethernet_connected_chip_ids()) {
+        for (const auto& [connected_device_id, cores] : device->get_ethernet_cores_grouped_by_connected_chips()) {
             for (uint32_t j = 0; j < devices.size(); ++j) {
                 if (devices[j]->id() == connected_device_id) {
                     adj[i][j] = 1;

tt_metal/impl/device/device.hpp

@@ -101,8 +101,8 @@ class Device {
 
     std::vector<CoreCoord> ethernet_cores_from_logical_cores(const std::vector<CoreCoord> &logical_cores) const;
 
-    std::unordered_set<chip_id_t> get_ethernet_connected_chip_ids() const {
-        return tt::Cluster::instance().get_ethernet_connected_chip_ids(this->id_);
+    std::unordered_map<chip_id_t, std::vector<CoreCoord>> get_ethernet_cores_grouped_by_connected_chips() const {
+        return tt::Cluster::instance().get_ethernet_cores_grouped_by_connected_chips(this->id_);
     }
 
     std::unordered_set<CoreCoord> get_active_ethernet_cores() const {
@@ -117,6 +117,10 @@ class Device {
         return tt::Cluster::instance().get_connected_ethernet_core(std::make_tuple(this->id_, eth_core));
     }
 
+    std::vector<CoreCoord> get_ethernet_sockets(chip_id_t connected_chip_id) const {
+        return tt::Cluster::instance().get_ethernet_sockets(this->id_, connected_chip_id);
+    }
+
     bool is_mmio_capable() const {
         return tt::Cluster::instance().get_associated_mmio_device(this->id_) == this->id_;
     }

tt_metal/llrt/tt_cluster.cpp

@@ -46,6 +46,8 @@ Cluster::Cluster() {
 
     this->initialize_device_drivers();
 
+    this->initialize_ethernet_sockets();
+
     this->assert_risc_reset();
 }
 
@@ -562,27 +564,69 @@ uint64_t Cluster::get_pcie_base_addr_from_device(chip_id_t chip_id) const {
 }
 
 // Ethernet cluster api
-std::unordered_set<chip_id_t> Cluster::get_ethernet_connected_chip_ids(chip_id_t chip_id) const {
-    std::unordered_set<chip_id_t> connected_chips;
+
+void Cluster::initialize_ethernet_sockets() {
+    for (const auto &chip_id : this->cluster_desc_->get_all_chips()) {
+        if (this->ethernet_sockets_.find(chip_id) == this->ethernet_sockets_.end()) {
+            this->ethernet_sockets_.insert({chip_id, {}});
+        }
+        for (const auto &[connected_chip_id, eth_cores] :
+             this->get_ethernet_cores_grouped_by_connected_chips(chip_id)) {
+            if (this->ethernet_sockets_.at(chip_id).find(connected_chip_id) ==
+                this->ethernet_sockets_.at(chip_id).end()) {
+                this->ethernet_sockets_.at(chip_id).insert({connected_chip_id, {}});
+            }
+            if (this->ethernet_sockets_.find(connected_chip_id) == this->ethernet_sockets_.end()) {
+                this->ethernet_sockets_.insert({connected_chip_id, {}});
+            }
+            if (this->ethernet_sockets_.at(connected_chip_id).find(chip_id) ==
+                this->ethernet_sockets_.at(connected_chip_id).end()) {
+                this->ethernet_sockets_.at(connected_chip_id).insert({chip_id, {}});
+            } else {
+                continue;
+            }
+            for (const auto &eth_core : eth_cores) {
+                this->ethernet_sockets_.at(chip_id).at(connected_chip_id).emplace_back(eth_core);
+                this->ethernet_sockets_.at(connected_chip_id)
+                    .at(chip_id)
+                    .emplace_back(std::get<1>(this->get_connected_ethernet_core(std::make_tuple(chip_id, eth_core))));
+            }
+        }
+    }
+}
+
+std::unordered_map<chip_id_t, std::vector<CoreCoord>> Cluster::get_ethernet_cores_grouped_by_connected_chips(
+    chip_id_t chip_id) const {
+    const auto &soc_desc = get_soc_desc(chip_id);
+    std::unordered_map<chip_id_t, std::vector<CoreCoord>> connected_chips;
     const auto &all_eth_connections = this->cluster_desc_->get_ethernet_connections();
     if (all_eth_connections.find(chip_id) == all_eth_connections.end()) {
         return {};
     }
     for (const auto &[eth_chan, connected_chip_chan] : all_eth_connections.at(chip_id)) {
-        connected_chips.insert(std::get<0>(connected_chip_chan));
+        const auto &other_chip_id = std::get<0>(connected_chip_chan);
+        if (connected_chips.find(other_chip_id) == connected_chips.end()) {
+            std::vector<CoreCoord> active_ethernet_cores;
+
+            for (const auto &channel_pair :
+                 this->cluster_desc_->get_directly_connected_ethernet_channels_between_chips(chip_id, other_chip_id)) {
+                ethernet_channel_t local_chip_chan = std::get<0>(channel_pair);
+                active_ethernet_cores.emplace_back(
+                    get_soc_desc(chip_id).chan_to_logical_eth_core_map.at(local_chip_chan));
+            }
+            connected_chips.insert({other_chip_id, active_ethernet_cores});
+        } else {
+            continue;
+        }
     }
     return connected_chips;
 }
 
 std::unordered_set<CoreCoord> Cluster::get_active_ethernet_cores(chip_id_t chip_id) const {
     std::unordered_set<CoreCoord> active_ethernet_cores;
-    const auto &connected_chips = this->get_ethernet_connected_chip_ids(chip_id);
-    for (auto &other_chip_id : connected_chips) {
-        for (const auto &channel_pair :
-             this->cluster_desc_->get_directly_connected_ethernet_channels_between_chips(chip_id, other_chip_id)) {
-            ethernet_channel_t local_chip_chan = std::get<0>(channel_pair);
-            active_ethernet_cores.insert(get_soc_desc(chip_id).chan_to_logical_eth_core_map.at(local_chip_chan));
-        }
+    const auto &connected_chips = this->get_ethernet_cores_grouped_by_connected_chips(chip_id);
+    for (const auto &[other_chip_id, eth_cores] : connected_chips) {
+        active_ethernet_cores.insert(eth_cores.begin(), eth_cores.end());
     }
     return active_ethernet_cores;
 }
@@ -612,6 +656,16 @@ std::tuple<chip_id_t, CoreCoord> Cluster::get_connected_ethernet_core(std::tuple
         std::get<0>(connected_eth_core), soc_desc.chan_to_logical_eth_core_map.at(std::get<1>(connected_eth_core)));
 }
 
+std::vector<CoreCoord> Cluster::get_ethernet_sockets(chip_id_t local_chip, chip_id_t remote_chip) const {
+    const auto &local_ethernet_sockets = this->ethernet_sockets_.at(local_chip);
+    TT_ASSERT(
+        local_ethernet_sockets.find(remote_chip) != local_ethernet_sockets.end(),
+        "Device {} is not connected to Device {}",
+        local_chip,
+        remote_chip);
+    return local_ethernet_sockets.at(remote_chip);
+}
+
 uint32_t Cluster::get_tensix_soft_reset_addr() const {
     return DEVICE_DATA.TENSIX_SOFT_RESET_ADDR;
 }

tt_metal/llrt/tt_cluster.hpp

@@ -107,8 +107,9 @@ class Cluster {
     uint64_t get_pcie_base_addr_from_device(chip_id_t chip_id) const;
 
     // Ethernet cluster api
-    // Returns set of connected chip ids
-    std::unordered_set<chip_id_t> get_ethernet_connected_chip_ids(chip_id_t chip_id) const;
+    // Returns map of connected chip ids to active ethernet cores
+    std::unordered_map<chip_id_t, std::vector<CoreCoord>> get_ethernet_cores_grouped_by_connected_chips(
+        chip_id_t chip_id) const;
 
     // Returns set of logical active ethernet coordinates on chip
     std::unordered_set<CoreCoord> get_active_ethernet_cores(chip_id_t chip_id) const;
@@ -119,6 +120,10 @@ class Cluster {
     // Returns connected ethernet core on the other chip
     std::tuple<chip_id_t, CoreCoord> get_connected_ethernet_core(std::tuple<chip_id_t, CoreCoord> eth_core) const;
 
+    // Returns a ethernet sockets between local chip and remote chip
+    // get_ethernet_sockets(a, b)[0] is connected to get_ethernet_sockets(b, a)[0]
+    std::vector<CoreCoord> get_ethernet_sockets(chip_id_t local_chip, chip_id_t remote_chip) const;
+
     // Returns MMIO device ID (logical) that controls given `device_id`. If `device_id` is MMIO device it is returned.
     chip_id_t get_associated_mmio_device(chip_id_t device_id) const {
         return this->device_to_mmio_device_.at(device_id);
@@ -153,6 +158,7 @@ class Cluster {
     tt_cxy_pair convert_physical_cxy_to_virtual(const tt_cxy_pair &physical_cxy) const;
     void configure_static_tlbs(chip_id_t mmio_device_id) const;
 
+    void initialize_ethernet_sockets();
 
     ARCH arch_;
     TargetDevice target_type_;
@@ -224,6 +230,8 @@ class Cluster {
         REQUEST_ROUTING_CMD_QUEUE_BASE,
         RESPONSE_ROUTING_CMD_QUEUE_BASE,
         CMD_BUF_PTR_MASK};
+
+    std::unordered_map<chip_id_t, std::unordered_map<chip_id_t, std::vector<CoreCoord>>> ethernet_sockets_;
 };
 
 }  // namespace tt