Improve and refactor harvesting and coordinate management

- Coordinate manager class for each arch

- Write tests for the harvesting

- Add functions to convert between coordinate systems

device/CMakeLists.txt

@@ -1,4 +1,3 @@
-
 set(POSITION_INDEPENDENT_CODE ON)
 
 set(UMD_DEVICE_SRCS
@@ -15,6 +14,9 @@ set(UMD_DEVICE_SRCS
     blackhole/blackhole_implementation.cpp
     grayskull/grayskull_implementation.cpp
     wormhole/wormhole_implementation.cpp
+    coordinate_manager.cpp
+    blackhole/blackhole_coordinate_manager.cpp
+    wormhole/wormhole_coordinate_manager.cpp
     pcie/pci_device.cpp
 )
 

device/blackhole/blackhole_coordinate_manager.cpp

@@ -0,0 +1,14 @@
+#include "blackhole_coordinate_manager.h"
+
+std::set<std::size_t> BlackholeCoordinateManager::get_x_coordinates_to_harvest(std::size_t harvesting_mask) {
+    std::set<std::size_t> x_to_harvest;
+    std::size_t logical_x = 0;
+    while (harvesting_mask > 0) {
+        if (harvesting_mask & 1) {
+            x_to_harvest.insert(logical_x);
+        }
+        logical_x++;
+        harvesting_mask >>= 1;
+    }
+    return x_to_harvest;
+}

device/blackhole/blackhole_coordinate_manager.h

@@ -0,0 +1,26 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2023 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <tuple>
+#include <vector>
+#include <map>
+#include <set>
+
+#include "device/coordinate_manager.h"
+
+class BlackholeCoordinateManager : public CoordinateManager {
+
+public:
+    BlackholeCoordinateManager(const tt_xy_pair& worker_grid_size, const std::vector<tt_xy_pair>& workers, std::size_t harvesting_mask)
+        : CoordinateManager(worker_grid_size, workers, harvesting_mask) {}
+
+protected: 
+    std::set<std::size_t> get_x_coordinates_to_harvest(std::size_t harvesting_mask) override;
+};

device/coordinate_manager.cpp

@@ -0,0 +1,161 @@
+#include "device/coordinate_manager.h"
+#include <memory>
+#include "coordinate_manager.h"
+#include "grayskull/grayskull_coordinate_manager.h"
+
+tt_physical_coords CoordinateManager::logical_to_physical_coords(tt_logical_coords logical_coords) {
+    // log_assert(logical_coords.x < logical_x_to_physical_x.size());
+    // log_assert(logical_coords.y < logical_y_to_physical_y.size());
+    return tt_physical_coords(logical_x_to_physical_x[logical_coords.x], logical_y_to_physical_y[logical_coords.y]);
+}
+
+// TODO(pjanevski): is it enough just to add 18 to the logical coordinates
+// in order to get the translated coordinates?
+tt_translated_coords CoordinateManager::logical_to_translated_coords(tt_logical_coords logical_coords) {
+    static const std::size_t translated_offset = 18;
+    return tt_translated_coords(logical_coords.x + translated_offset, logical_coords.y + translated_offset);
+}
+
+tt_logical_coords CoordinateManager::physical_to_logical_coords(tt_physical_coords physical_coords) {
+    return tt_logical_coords(physical_x_to_logical_x[physical_coords.x], physical_y_to_logical_y[physical_coords.y]);
+}
+
+tt_translated_coords CoordinateManager::physical_to_translated_coords(tt_physical_coords physical_coords) {
+    return tt_translated_coords(0, 0);
+}
+
+tt_virtual_coords CoordinateManager::logical_to_virtual_coords(tt_logical_coords logical_coords) {
+    return tt_virtual_coords(logical_x_to_virtual_x[logical_coords.x], logical_y_to_virtual_y[logical_coords.y]);
+}
+
+tt_logical_coords CoordinateManager::virtual_to_logical_coords(tt_virtual_coords virtual_coords) {
+    return tt_logical_coords(virtual_x_to_logical_x[virtual_coords.x], virtual_y_to_logical_y[virtual_coords.y]);
+}
+
+void CoordinateManager::clear_harvesting_structures() {
+    logical_x_to_physical_x.clear();
+    logical_y_to_physical_y.clear();
+    logical_x_to_virtual_x.clear();
+    logical_y_to_virtual_y.clear();
+    physical_x_to_logical_x.clear();
+    physical_y_to_logical_y.clear();
+    virtual_x_to_logical_x.clear();
+    virtual_y_to_logical_y.clear();
+}
+
+std::set<std::size_t> CoordinateManager::get_x_coordinates_to_harvest(std::size_t harvesting_mask) {
+    return {};
+}
+
+std::set<std::size_t> CoordinateManager::get_y_coordinates_to_harvest(std::size_t harvesting_mask) {
+    return {};
+}
+
+void CoordinateManager::perform_harvesting(std::size_t harvesting_mask) {
+    clear_harvesting_structures();
+
+    std::set<size_t> physical_x_unharvested;
+    std::set<size_t> physical_y_unharvested;
+    for (auto core : workers) {
+        physical_x_unharvested.insert(core.x);
+        physical_y_unharvested.insert(core.y);
+    }
+
+    std::set<std::size_t> x_coordinates_to_harvest = get_x_coordinates_to_harvest(harvesting_mask);
+    std::set<std::size_t> y_coordinates_to_harvest = get_y_coordinates_to_harvest(harvesting_mask);
+
+    std::size_t num_harvested_y = y_coordinates_to_harvest.size();
+    std::size_t num_harvested_x = x_coordinates_to_harvest.size();
+
+    std::size_t grid_size_x = worker_grid_size.x;
+    std::size_t grid_size_y = worker_grid_size.y;
+
+    logical_x_to_physical_x.resize(grid_size_x - num_harvested_x);
+    logical_y_to_physical_y.resize(grid_size_y - num_harvested_y);
+
+    logical_x_to_virtual_x.resize(grid_size_x - num_harvested_x);
+    logical_y_to_virtual_y.resize(grid_size_y - num_harvested_y);
+
+    fill_logical_to_physical_mapping(x_coordinates_to_harvest, y_coordinates_to_harvest, physical_x_unharvested, physical_y_unharvested);
+    fill_logical_to_virtual_mapping(physical_x_unharvested, physical_y_unharvested);
+}
+
+void CoordinateManager::fill_logical_to_physical_mapping(
+    const std::set<size_t>& x_to_harvest, const std::set<size_t>& y_to_harvest,
+    const std::set<size_t>& physical_x_unharvested, const std::set<size_t>& physical_y_unharvested) {
+
+    auto physical_y_it = physical_y_unharvested.begin();
+    std::size_t logical_y = 0;
+    for (size_t y = 0; y < worker_grid_size.y; y++) {
+        if (y_to_harvest.find(y) == y_to_harvest.end()) {
+            logical_y_to_physical_y[logical_y] = *physical_y_it;
+            if (physical_y_to_logical_y.find(*physical_y_it) != physical_y_to_logical_y.end()) {
+                throw std::runtime_error("Duplicate physical y coordinate found in the worker cores");
+            }
+            physical_y_to_logical_y[*physical_y_it] = logical_y;
+            logical_y++;
+            physical_y_it++;
+        } else {
+            physical_y_it++;
+        }
+    }
+
+    auto physical_x_it = physical_x_unharvested.begin();
+    std::size_t logical_x = 0;
+    for(std::size_t x = 0; x < worker_grid_size.x; x++) {
+        if (x_to_harvest.find(x) == x_to_harvest.end()) {
+            logical_x_to_physical_x[logical_x] = *physical_x_it;
+            if (physical_x_to_logical_x.find(*physical_x_it) != physical_x_to_logical_x.end()) {
+                throw std::runtime_error("Duplicate physical x coordinate found in the worker cores");
+            }
+            physical_x_to_logical_x[*physical_x_it] = logical_x;
+            logical_x++;
+            physical_x_it++;
+        } else {
+            physical_x_it++;
+        }
+    }
+}
+
+void CoordinateManager::fill_logical_to_virtual_mapping(const std::set<size_t>& physical_x_unharvested, const std::set<size_t>& physical_y_unharvested) {
+    auto physical_y_it = physical_y_unharvested.begin();
+    for (std::size_t y = 0; y < logical_y_to_virtual_y.size(); y++) {
+        logical_y_to_virtual_y[y] = *physical_y_it;
+        if (virtual_y_to_logical_y.find(*physical_y_it) != virtual_y_to_logical_y.end()) {
+            throw std::runtime_error("Duplicate virtual y coordinate found in the worker cores");
+        }
+        virtual_y_to_logical_y[*physical_y_it] = y;
+        physical_y_it++;
+    }
+
+    auto physical_x_it = physical_x_unharvested.begin();
+    for (std::size_t x = 0; x < logical_x_to_virtual_x.size(); x++) {
+        logical_x_to_virtual_x[x] = *physical_x_it;
+        if (virtual_x_to_logical_x.find(*physical_x_it) != virtual_x_to_logical_x.end()) {
+            throw std::runtime_error("Duplicate virtual x coordinate found in the worker cores");
+        }
+        virtual_x_to_logical_x[*physical_x_it] = x;
+        physical_x_it++;
+    }
+}
+
+#include "device/blackhole/blackhole_coordinate_manager.h"
+#include "device/grayskull/grayskull_coordinate_manager.h"
+#include "device/wormhole/wormhole_coordinate_manager.h"
+
+CoordinateManager* CoordinateManager::get_coordinate_manager(
+    tt::ARCH arch,
+    const tt_xy_pair& worker_grid_size,
+    const std::vector<tt_xy_pair>& workers,
+    std::size_t harvesting_mask) {
+
+    if (arch == tt::ARCH::GRAYSKULL) {
+        return new GrayskullCoordinateManager(worker_grid_size, workers, harvesting_mask);
+    } else if (arch == tt::ARCH::WORMHOLE_B0) {
+        return new WormholeCoordinateManager(worker_grid_size, workers, harvesting_mask);
+    } else if (arch == tt::ARCH::BLACKHOLE) {
+        return new BlackholeCoordinateManager(worker_grid_size, workers, harvesting_mask);
+    } else {
+        throw std::runtime_error("Invalid architecture for coordinate manager");
+    }
+}

device/coordinate_manager.h

@@ -0,0 +1,67 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2023 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <tuple>
+#include <vector>
+#include <map>
+#include <set>
+
+#include "device/tt_xy_pair.h"
+#include "device/tt_arch_types.h"
+#include <iostream>
+
+class CoordinateManager {
+
+public:
+    CoordinateManager(const tt_xy_pair& worker_grid_size, const std::vector<tt_xy_pair>& workers, std::size_t harvesting_mask)
+        : worker_grid_size(worker_grid_size), workers(workers), harvesting_mask(harvesting_mask) {}
+
+    virtual void perform_harvesting(std::size_t harvesting_mask);
+
+    virtual tt_physical_coords logical_to_physical_coords(tt_logical_coords logical_coords);
+    virtual tt_translated_coords logical_to_translated_coords(tt_logical_coords logical_coords);
+    virtual tt_logical_coords physical_to_logical_coords(tt_physical_coords physical_coords);
+    virtual tt_translated_coords physical_to_translated_coords(tt_physical_coords physical_coords);
+    virtual tt_virtual_coords logical_to_virtual_coords(tt_logical_coords logical_coords);
+    virtual tt_logical_coords virtual_to_logical_coords(tt_virtual_coords virtual_coords);
+
+    static CoordinateManager* get_coordinate_manager(
+        tt::ARCH arch,
+        const tt_xy_pair& worker_grid_size,
+        const std::vector<tt_xy_pair>& workers,
+        std::size_t harvesting_mask);
+
+protected:
+    virtual void clear_harvesting_structures();
+
+    virtual std::set<std::size_t> get_x_coordinates_to_harvest(std::size_t harvesting_mask);
+    virtual std::set<std::size_t> get_y_coordinates_to_harvest(std::size_t harvesting_mask);
+
+    virtual void fill_logical_to_physical_mapping(
+        const std::set<size_t>& x_to_harvest, const std::set<size_t>& y_to_harvest,
+        const std::set<size_t>& physical_x_unharvested, const std::set<size_t>& physical_y_unharvested);
+    virtual void fill_logical_to_virtual_mapping(const std::set<size_t>& physical_x_unharvested, const std::set<size_t>& physical_y_unharvested);
+
+    std::map<std::size_t, std::size_t> physical_y_to_logical_y;
+    std::map<std::size_t, std::size_t> physical_x_to_logical_x;
+
+    std::vector<std::size_t> logical_y_to_physical_y;
+    std::vector<std::size_t> logical_x_to_physical_x;
+
+    std::vector<std::size_t> logical_y_to_virtual_y;
+    std::vector<std::size_t> logical_x_to_virtual_x;
+
+    std::map<std::size_t, std::size_t> virtual_y_to_logical_y;
+    std::map<std::size_t, std::size_t> virtual_x_to_logical_x;
+
+    const tt_xy_pair worker_grid_size;
+    const std::vector<tt_xy_pair>& workers;
+    const std::size_t harvesting_mask;
+};

device/grayskull/grayskull_coordinate_manager.h

@@ -0,0 +1,23 @@
+/*
+ * SPDX-FileCopyrightText: (c) 2023 Tenstorrent Inc.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include <tuple>
+#include <vector>
+#include <map>
+#include <set>
+
+#include "device/coordinate_manager.h"
+
+class GrayskullCoordinateManager : public CoordinateManager {
+
+public:
+    GrayskullCoordinateManager(const tt_xy_pair& worker_grid_size, const std::vector<tt_xy_pair>& workers, std::size_t harvesting_mask)
+        : CoordinateManager(worker_grid_size, workers, harvesting_mask) {}
+};

device/tt_soc_descriptor.cpp

@@ -2,13 +2,16 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#include "coordinate_manager.h"
+#include "tt_xy_pair.h"
 #include "yaml-cpp/yaml.h"
 #include "tt_soc_descriptor.h"
 
 #include <assert.h>
 #include <fstream>
 #include <iostream>
 #include <regex>
+#include <stdexcept>
 #include <string>
 #include <unordered_set>
 
@@ -58,10 +61,7 @@ inline std::string& trim(std::string& s, const char* t = ws)
 }
 
 void tt_SocDescriptor::load_soc_features_from_device_descriptor(YAML::Node &device_descriptor_yaml) {
-    overlay_version = device_descriptor_yaml["features"]["overlay"]["version"].as<int>();
     noc_translation_id_enabled = device_descriptor_yaml["features"]["noc"] && device_descriptor_yaml["features"]["noc"]["translation_id_enabled"] ? device_descriptor_yaml["features"]["noc"]["translation_id_enabled"].as<bool>() : false;
-    packer_version = device_descriptor_yaml["features"]["packer"]["version"].as<int>();
-    unpacker_version = device_descriptor_yaml["features"]["unpacker"]["version"].as<int>();
     dst_size_alignment = device_descriptor_yaml["features"]["math"]["dst_size_alignment"].as<int>();
     worker_l1_size = device_descriptor_yaml["worker_l1_size"].as<int>();
     eth_l1_size = device_descriptor_yaml["eth_l1_size"].as<int>();
@@ -166,7 +166,39 @@ void tt_SocDescriptor::load_core_descriptors_from_device_descriptor(YAML::Node &
     }
 }
 
-tt_SocDescriptor::tt_SocDescriptor(std::string device_descriptor_path) {
+void tt_SocDescriptor::create_coordinate_manager(std::size_t harvesting_mask) {
+    coordinate_manager = CoordinateManager::get_coordinate_manager(arch, worker_grid_size, workers, harvesting_mask);
+}
+
+void tt_SocDescriptor::perform_harvesting(std::size_t harvesting_mask) {
+    coordinate_manager->perform_harvesting(harvesting_mask);
+}
+
+tt_physical_coords tt_SocDescriptor::logical_to_physical_coords(tt_logical_coords logical_coords) {
+    return coordinate_manager->logical_to_physical_coords(logical_coords);
+}
+
+tt_translated_coords tt_SocDescriptor::logical_to_translated_coords(tt_logical_coords logical_coords) {
+    return coordinate_manager->logical_to_translated_coords(logical_coords);
+}
+
+tt_logical_coords tt_SocDescriptor::physical_to_logical_coords(tt_physical_coords physical_coords) {
+    return coordinate_manager->physical_to_logical_coords(physical_coords);
+}
+
+tt_translated_coords tt_SocDescriptor::physical_to_translated_coords(tt_physical_coords physical_coords) {
+    return coordinate_manager->physical_to_translated_coords(physical_coords);
+}
+
+tt_virtual_coords tt_SocDescriptor::logical_to_virtual_coords(tt_logical_coords logical_coords) {
+    return coordinate_manager->logical_to_virtual_coords(logical_coords);
+}
+
+tt_logical_coords tt_SocDescriptor::virtual_to_logical_coords(tt_virtual_coords virtual_coords) {
+    return coordinate_manager->virtual_to_logical_coords(virtual_coords);
+}
+
+tt_SocDescriptor::tt_SocDescriptor(std::string device_descriptor_path, std::size_t harvesting_mask) {
     std::ifstream fdesc(device_descriptor_path);
     if (fdesc.fail()) {
         throw std::runtime_error(fmt::format("Error: device descriptor file {} does not exist!", device_descriptor_path));
@@ -189,6 +221,8 @@ tt_SocDescriptor::tt_SocDescriptor(std::string device_descriptor_path) {
     arch_name_value = trim(arch_name_value);
     arch = get_arch_name(arch_name_value);
     load_soc_features_from_device_descriptor(device_descriptor_yaml);
+    create_coordinate_manager(harvesting_mask);
+    perform_harvesting(harvesting_mask);
 }
 
 int tt_SocDescriptor::get_num_dram_channels() const {