Initial Chip class (#351)

### Issue
First PR towards #248 

### Description
First in set of changes to introduce Chip abstraction class according to
the plan
https://docs.google.com/drawings/d/1-m1azdsBqMA0A6ATYRMfkhyeuOJuGCEI62N5a96LXj0/edit
This change also required minor rethinking of some parts of how Cluster
is constructed (this took much time to figure out).

### List of the changes
- Added Chip class, and derived Local, Remote, and Mock classes. Local
will be created for mmio chips, Remote for non-mmio. Mock chip is now
only a placeholder, to show the Cluster constructor API changes, but it
should merge with existing MockupDevice. Chip is an abstract class.
- Chip classes for now only hold soc descriptors and nothing else.
Further PRs will gradually move out logic from Cluster to Chip.
- get_soc_descriptor and soc_descriptor_per_chip were moved around to
conform to the new reality
- This was reverted after offline discussion: (Cluster constructor which
takes soc descriptor was changed so that it takes tt_SocDescriptor
object rather than yaml file. This is because the object also holds
harvesting information.)
- Added another Cluster constructor which takes a set of chips. This is
an experimental one, and users should know what they're doing and they
can break the code easily.
- Previously mentioned new constructor was used to offer
Cluster::create_mock_cluster which is a helper function to build you a
cluster with mock chip.
- Wrote a test which showcases different constructor usages, which was a
leftover from #277

### Testing
Existing tests should prove no functional changes in this PR.

### API Changes
This PR has no API changes.

device/CMakeLists.txt

@@ -24,6 +24,10 @@ target_sources(
     device
     PRIVATE
         architecture_implementation.cpp
+        chip/chip.cpp
+        chip/local_chip.cpp
+        chip/mock_chip.cpp
+        chip/remote_chip.cpp
         cluster.cpp
         coordinate_manager.cpp
         cpuset_lib.cpp

device/api/umd/device/chip/chip.h

@@ -0,0 +1,29 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "umd/device/tt_soc_descriptor.h"
+#include "umd/device/types/cluster_descriptor_types.h"
+
+namespace tt::umd {
+
+// An abstract class that represents a chip.
+class Chip {
+public:
+    Chip(tt_SocDescriptor soc_descriptor);
+
+    virtual ~Chip() = default;
+
+    tt_SocDescriptor& get_soc_descriptor();
+
+    virtual bool is_mmio_capable() const = 0;
+
+private:
+    tt_SocDescriptor soc_descriptor_;
+};
+
+}  // namespace tt::umd

device/api/umd/device/chip/local_chip.h

@@ -0,0 +1,17 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "umd/device/chip/chip.h"
+
+namespace tt::umd {
+class LocalChip : public Chip {
+public:
+    LocalChip(tt_SocDescriptor soc_descriptor);
+    bool is_mmio_capable() const override;
+};
+}  // namespace tt::umd

device/api/umd/device/chip/mock_chip.h

@@ -0,0 +1,17 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "umd/device/chip/chip.h"
+
+namespace tt::umd {
+class MockChip : public Chip {
+public:
+    MockChip(tt_SocDescriptor soc_descriptor);
+    bool is_mmio_capable() const override;
+};
+}  // namespace tt::umd

device/api/umd/device/chip/remote_chip.h

@@ -0,0 +1,17 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "umd/device/chip/chip.h"
+
+namespace tt::umd {
+class RemoteChip : public Chip {
+public:
+    RemoteChip(tt_SocDescriptor soc_descriptor);
+    bool is_mmio_capable() const override;
+};
+}  // namespace tt::umd

device/api/umd/device/cluster.h

@@ -17,6 +17,7 @@
 #include "tt_silicon_driver_common.hpp"
 #include "tt_soc_descriptor.h"
 #include "tt_xy_pair.h"
+#include "umd/device/chip/chip.h"
 #include "umd/device/tlb.h"
 #include "umd/device/tt_device/tt_device.h"
 #include "umd/device/tt_io.hpp"
@@ -41,8 +42,8 @@ class tt_ClusterDescriptor;
  */
 class tt_device {
 public:
-    tt_device();
-    virtual ~tt_device();
+    tt_device(){};
+    virtual ~tt_device(){};
 
     // Setup/Teardown Functions
     /**
@@ -279,8 +280,8 @@ class tt_device {
      * Query post harvesting SOC descriptors from UMD in virtual coordinates.
      * These descriptors should be used for looking up cores that are passed into UMD APIs.
      */
-    virtual std::unordered_map<chip_id_t, tt_SocDescriptor>& get_virtual_soc_descriptors() {
-        throw std::runtime_error("---- tt_device:get_virtual_soc_descriptors is not implemented\n");
+    virtual std::unordered_map<chip_id_t, tt_SocDescriptor> get_virtual_soc_descriptors() {
+        return soc_descriptor_per_chip;
     }
 
     /**
@@ -448,13 +449,16 @@ class tt_device {
         return 0;
     }
 
-    const tt_SocDescriptor& get_soc_descriptor(chip_id_t chip_id) const;
+    virtual const tt_SocDescriptor& get_soc_descriptor(chip_id_t chip_id) const {
+        return soc_descriptor_per_chip.at(chip_id);
+    }
 
     bool performed_harvesting = false;
     std::unordered_map<chip_id_t, uint32_t> harvested_rows_per_target = {};
     bool translation_tables_en = false;
 
 protected:
+    // TODO: Remove this once get_virtual_soc_descriptors can be removed.
     std::unordered_map<chip_id_t, tt_SocDescriptor> soc_descriptor_per_chip = {};
 };
 
@@ -469,17 +473,35 @@ class Cluster : public tt_device {
     // Constructor
     /**
      * Cluster constructor.
+     * Simplest form, creates a cluster of all available devices on the system.
+     *
+     * @param num_host_mem_ch_per_mmio_device Requested number of host channels (hugepages).
+     * @param skip_driver_allocs
+     * @param clean_system_resource Specifies if host state from previous runs needs to be cleaned up.
+     * @param perform_harvesting Allow the driver to modify the SOC descriptors per chip.
+     * @param simulated_harvesting_masks Manually specify additional harvesting masks for the devices in the cluster.
+     * The ones defined by the devices itself have to be used, they will be merged with the ones passed here.
+     */
+    Cluster(
+        const uint32_t& num_host_mem_ch_per_mmio_device = 1,
+        const bool skip_driver_allocs = false,
+        const bool clean_system_resources = false,
+        bool perform_harvesting = true,
+        std::unordered_map<chip_id_t, uint32_t> simulated_harvesting_masks = {});
+
+    /**
+     * Cluster constructor.
+     * This constructor can be used to target only specific devices on the system.
      *
-     * @param sdesc_path SOC descriptor specifying single chip.
      * @param target_devices Devices to target.
      * @param num_host_mem_ch_per_mmio_device Requested number of host channels (hugepages).
      * @param skip_driver_allocs
      * @param clean_system_resource Specifies if host state from previous runs needs to be cleaned up.
      * @param perform_harvesting Allow the driver to modify the SOC descriptors per chip.
-     * @param simulated_harvesting_masks
+     * @param simulated_harvesting_masks Manually specify additional harvesting masks for the devices in the cluster.
+     * The ones defined by the devices itself have to be used, they will be merged with the ones passed here.
      */
     Cluster(
-        const std::string& sdesc_path,
         const std::set<chip_id_t>& target_devices,
         const uint32_t& num_host_mem_ch_per_mmio_device = 1,
         const bool skip_driver_allocs = false,
@@ -488,42 +510,57 @@ class Cluster : public tt_device {
         std::unordered_map<chip_id_t, uint32_t> simulated_harvesting_masks = {});
 
     /**
-     * Cluster constructor. This constructor should be used to work towards removing all
-     * of the params from the constructor of tt_SiliconDevice (to become Cluster).
+     * Cluster constructor.
+     * This constructor can be used with custom soc descriptors for the devices on the system.
      *
+     * @param sdesc_path SOC descriptor yaml path specifying single chip. The passed soc descriptor will be used as a
+     * default device description for devices in the cluster, but each chip will be harvested according to the
+     * harvesting info of the devices in the cluster.
+     * @param target_devices Devices to target.
      * @param num_host_mem_ch_per_mmio_device Requested number of host channels (hugepages).
      * @param skip_driver_allocs
      * @param clean_system_resource Specifies if host state from previous runs needs to be cleaned up.
      * @param perform_harvesting Allow the driver to modify the SOC descriptors per chip.
-     * @param simulated_harvesting_masks
+     * @param simulated_harvesting_masks Manually specify additional harvesting masks for the devices in the cluster.
+     * The ones defined by the devices itself have to be used, they will be merged with the ones passed here.
      */
     Cluster(
+        const std::string& sdesc_path,
+        const std::set<chip_id_t>& target_devices,
         const uint32_t& num_host_mem_ch_per_mmio_device = 1,
         const bool skip_driver_allocs = false,
         const bool clean_system_resources = false,
         bool perform_harvesting = true,
         std::unordered_map<chip_id_t, uint32_t> simulated_harvesting_masks = {});
 
     /**
-     * Cluster constructor. This constructor should be used to target specific devices in a cluster.
+     * Cluster constructor.
+     * This constructor offers maximal flexibility, allowing the user to pass manually created Chips.
+     * The user has to know what they are doing.
+     * TODO: Could fail if logical_ids not match the ones in cluster descriptor, while Cluster still uses cluster
+     * descriptor.
      *
-     * @param target_devices Devices to target.
+     * @param chips Map of logical device ids to Chip instances.
      * @param num_host_mem_ch_per_mmio_device Requested number of host channels (hugepages).
      * @param skip_driver_allocs
      * @param clean_system_resource Specifies if host state from previous runs needs to be cleaned up.
      * @param perform_harvesting Allow the driver to modify the SOC descriptors per chip.
      * @param simulated_harvesting_masks
      */
     Cluster(
-        const std::set<chip_id_t>& target_devices,
+        std::unordered_map<chip_id_t, std::unique_ptr<Chip>>& chips,
         const uint32_t& num_host_mem_ch_per_mmio_device = 1,
         const bool skip_driver_allocs = false,
         const bool clean_system_resources = false,
         bool perform_harvesting = true,
         std::unordered_map<chip_id_t, uint32_t> simulated_harvesting_masks = {});
 
+    /**
+     * Cluster constructor which creates a cluster with Mock chips.
+     */
+    static std::unique_ptr<Cluster> create_mock_cluster();
+
     // Setup/Teardown Functions
-    virtual std::unordered_map<chip_id_t, tt_SocDescriptor>& get_virtual_soc_descriptors();
     virtual void set_device_l1_address_params(const tt_device_l1_address_params& l1_address_params_);
     virtual void set_device_dram_address_params(const tt_device_dram_address_params& dram_address_params_);
     virtual void set_driver_host_address_params(const tt_driver_host_address_params& host_address_params_);
@@ -638,14 +675,18 @@ class Cluster : public tt_device {
 
     TTDevice* get_tt_device(chip_id_t device_id) const;
 
+    const tt_SocDescriptor& get_soc_descriptor(chip_id_t chip_id) const;
+    // TODO: This function should be removed.
+    std::unordered_map<chip_id_t, tt_SocDescriptor> get_virtual_soc_descriptors();
+
     // Destructor
     virtual ~Cluster();
 
 private:
     // Helper functions
     // Startup + teardown
     void create_device(
-        const std::unordered_set<chip_id_t>& target_mmio_device_ids,
+        const std::set<chip_id_t>& target_mmio_device_ids,
         const uint32_t& num_host_mem_ch_per_mmio_device,
         const bool skip_driver_allocs,
         const bool clean_system_resources);
@@ -656,7 +697,7 @@ class Cluster : public tt_device {
     void broadcast_tensix_risc_reset_to_cluster(const TensixSoftResetOptions& soft_resets);
     void send_remote_tensix_risc_reset_to_core(const tt_cxy_pair& core, const TensixSoftResetOptions& soft_resets);
     void send_tensix_risc_reset_to_core(const tt_cxy_pair& core, const TensixSoftResetOptions& soft_resets);
-    void perform_harvesting_and_populate_soc_descriptors(const std::string& sdesc_path, const bool perform_harvesting);
+    void perform_harvesting_on_soc_descriptors();
     void populate_cores();
     void init_pcie_iatus();  // No more p2p support.
     void check_pcie_device_initialized(int device_id);
@@ -771,9 +812,11 @@ class Cluster : public tt_device {
 
     // This functions has to be called for local chip, and then it will wait for all connected remote chips to flush.
     void wait_for_connected_non_mmio_flush(chip_id_t chip_id);
-
+    std::unique_ptr<Chip> construct_chip_from_cluster(
+        chip_id_t chip_id, tt_ClusterDescriptor* cluster_desc, tt_SocDescriptor& soc_desc);
+    std::unique_ptr<Chip> construct_chip_from_cluster(chip_id_t logical_device_id, tt_ClusterDescriptor* cluster_desc);
+    void add_chip(chip_id_t chip_id, std::unique_ptr<Chip> chip);
     void construct_cluster(
-        const std::string& sdesc_path,
         const uint32_t& num_host_mem_ch_per_mmio_device,
         const bool skip_driver_allocs,
         const bool clean_system_resources,
@@ -789,6 +832,8 @@ class Cluster : public tt_device {
     std::vector<tt::ARCH> archs_in_cluster = {};
     std::set<chip_id_t> target_devices_in_cluster = {};
     std::set<chip_id_t> target_remote_chips = {};
+    std::set<chip_id_t> all_target_mmio_devices = {};
+    std::unordered_map<chip_id_t, std::unique_ptr<Chip>> chips_;
     tt::ARCH arch_name;
 
     // Map of enabled tt devices
@@ -820,7 +865,6 @@ class Cluster : public tt_device {
     std::unordered_set<tt_xy_pair> eth_cores = {};
     std::unordered_set<tt_xy_pair> dram_cores = {};
     std::map<chip_id_t, std::unordered_map<int32_t, uint64_t>> tlb_config_map = {};
-    std::set<chip_id_t> all_target_mmio_devices;
 
     // Note that these maps holds only entries for local PCIe chips.
     std::unordered_map<chip_id_t, std::function<std::int32_t(tt_xy_pair)>> map_core_to_tlb_per_chip = {};

device/api/umd/device/tt_simulation_device.h

@@ -20,8 +20,6 @@ class tt_SimulationDevice : public tt_device {
 
     tt_SimulationHost host;
 
-    // Setup/Teardown Functions
-    virtual std::unordered_map<chip_id_t, tt_SocDescriptor>& get_virtual_soc_descriptors();
     virtual void set_device_l1_address_params(const tt_device_l1_address_params& l1_address_params_);
     virtual void set_device_dram_address_params(const tt_device_dram_address_params& dram_address_params_);
     virtual void set_driver_host_address_params(const tt_driver_host_address_params& host_address_params_);

device/chip/chip.cpp

@@ -0,0 +1,15 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "umd/device/chip/chip.h"
+
+namespace tt::umd {
+
+Chip::Chip(tt_SocDescriptor soc_descriptor) : soc_descriptor_(soc_descriptor) {}
+
+tt_SocDescriptor& Chip::get_soc_descriptor() { return soc_descriptor_; }
+
+}  // namespace tt::umd

device/chip/local_chip.cpp

@@ -0,0 +1,14 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "umd/device/chip/local_chip.h"
+
+namespace tt::umd {
+
+LocalChip::LocalChip(tt_SocDescriptor soc_descriptor) : Chip(soc_descriptor) {}
+
+bool LocalChip::is_mmio_capable() const { return true; }
+}  // namespace tt::umd

device/chip/mock_chip.cpp

@@ -0,0 +1,14 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "umd/device/chip/mock_chip.h"
+
+namespace tt::umd {
+
+MockChip::MockChip(tt_SocDescriptor soc_descriptor) : Chip(soc_descriptor) {}
+
+bool MockChip::is_mmio_capable() const { return true; }
+}  // namespace tt::umd

device/chip/remote_chip.cpp

@@ -0,0 +1,14 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2024 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "umd/device/chip/remote_chip.h"
+
+namespace tt::umd {
+
+RemoteChip::RemoteChip(tt_SocDescriptor soc_descriptor) : Chip(soc_descriptor) {}
+
+bool RemoteChip::is_mmio_capable() const { return false; }
+}  // namespace tt::umd