From 342feac8f3154271ba06640c4805caaa1636edff Mon Sep 17 00:00:00 2001
From: Bojan Rosko <brosko@tenstorrent.com>
Date: Tue, 26 Nov 2024 09:55:00 +0000
Subject: [PATCH] init

---
 common/disjoint_set.hpp               |  2 +-
 tests/api/test_cluster.cpp            | 18 ++++++++++--------
 tests/api/test_cluster_descriptor.cpp | 11 ++++++++++-
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/common/disjoint_set.hpp b/common/disjoint_set.hpp
index b2187173..82884380 100644
--- a/common/disjoint_set.hpp
+++ b/common/disjoint_set.hpp
@@ -24,7 +24,7 @@ class DisjointSet {
     void merge(T item1, T item2) {
         T set1 = get_set(item1);
         T set2 = get_set(item2);
-        parent[set1] = set2;
+        parent[set1] = parent[set2] = std::min(set1, set2);
     }
 
     bool are_same_set(T item1, T item2) { return get_set(item1) == get_set(item2); }
diff --git a/tests/api/test_cluster.cpp b/tests/api/test_cluster.cpp
index fa034506..9ed1589e 100644
--- a/tests/api/test_cluster.cpp
+++ b/tests/api/test_cluster.cpp
@@ -59,12 +59,12 @@ TEST(ApiClusterTest, OpenAllChips) { std::unique_ptr<Cluster> umd_cluster = get_
 TEST(ApiClusterTest, SimpleIOAllChips) {
     std::unique_ptr<Cluster> umd_cluster = get_cluster();
 
-    const tt_ClusterDescriptor* cluster_desc = umd_cluster->get_cluster_description();
-
     if (umd_cluster == nullptr || umd_cluster->get_all_chips_in_cluster().empty()) {
         GTEST_SKIP() << "No chips present on the system. Skipping test.";
     }
 
+    const tt_ClusterDescriptor* cluster_desc = umd_cluster->get_cluster_description();
+
     // Initialize random data.
     size_t data_size = 1024;
     std::vector<uint8_t> data(data_size, 0);
@@ -117,12 +117,12 @@ TEST(ApiClusterTest, SimpleIOAllChips) {
 TEST(ApiClusterTest, RemoteFlush) {
     std::unique_ptr<Cluster> umd_cluster = get_cluster();
 
-    const tt_ClusterDescriptor* cluster_desc = umd_cluster->get_cluster_description();
-
     if (umd_cluster == nullptr || umd_cluster->get_all_chips_in_cluster().empty()) {
         GTEST_SKIP() << "No chips present on the system. Skipping test.";
     }
 
+    const tt_ClusterDescriptor* cluster_desc = umd_cluster->get_cluster_description();
+
     size_t data_size = 1024;
     std::vector<uint8_t> data(data_size, 0);
 
@@ -175,14 +175,16 @@ TEST(ApiClusterTest, RemoteFlush) {
 }
 
 TEST(ApiClusterTest, SimpleIOSpecificChips) {
+    std::vector<int> pci_device_ids = PCIDevice::enumerate_devices();
+    // TODO: Make this test work on a host system without any tt devices.
+    if (pci_device_ids.empty()) {
+        GTEST_SKIP() << "No chips present on the system. Skipping test.";
+    }
+
     std::unique_ptr<Cluster> umd_cluster = std::make_unique<Cluster>(0);
 
     const tt_ClusterDescriptor* cluster_desc = umd_cluster->get_cluster_description();
 
-    if (umd_cluster == nullptr || umd_cluster->get_all_chips_in_cluster().empty()) {
-        GTEST_SKIP() << "No chips present on the system. Skipping test.";
-    }
-
     // Initialize random data.
     size_t data_size = 1024;
     std::vector<uint8_t> data(data_size, 0);
diff --git a/tests/api/test_cluster_descriptor.cpp b/tests/api/test_cluster_descriptor.cpp
index af14e77c..f4a6f136 100644
--- a/tests/api/test_cluster_descriptor.cpp
+++ b/tests/api/test_cluster_descriptor.cpp
@@ -99,6 +99,16 @@ TEST(ApiClusterDescriptorTest, TestAllOfflineClusterDescriptors) {
 
         std::unordered_map<chip_id_t, std::unordered_set<chip_id_t>> chips_grouped_by_closest_mmio =
             cluster_desc->get_chips_grouped_by_closest_mmio();
+
+        // Check that cluster_id is always the same for the same cluster.
+        // Cluster id takes the value of the smallest chip_id in the cluster.
+        for (auto const &[chip, coord] : eth_chip_coords) {
+            if (cluster_desc_yaml != "wormhole_2xN300_unconnected.yaml") {
+                EXPECT_EQ(coord.cluster_id, 0);
+            } else {
+                EXPECT_TRUE(coord.cluster_id == 0 || coord.cluster_id == 1);
+            }
+        }
     }
 }
 
@@ -125,7 +135,6 @@ TEST(ApiClusterDescriptorTest, SeparateClusters) {
     std::cout << "Detected " << chip_clusters.get_num_sets() << " separate clusters." << std::endl;
 
     // Check that get_closes_mmio_capable_chip works.
-    // Currently, it is expected that the following fails if there is more than 1 cluster.
     for (auto chip : all_chips) {
         chip_id_t closest_mmio_chip = cluster_desc->get_closest_mmio_capable_chip(chip);
         EXPECT_TRUE(chip_clusters.are_same_set(chip, closest_mmio_chip));