Consider register for mapping

include/NeuraDialect/Architecture/Architecture.h

@@ -172,6 +172,14 @@ class Tile : public BasicResource {
     return true;
   }
 
+  void addRegisterFileCluster(RegisterFileCluster* register_file_cluster);
+
+  const RegisterFileCluster *getRegisterFileCluster() const;
+
+  const std::vector<RegisterFile *> getRegisterFiles() const;
+
+  const std::vector<Register *> getRegisters() const;
+
 private:
   int id;
   int x, y;
@@ -181,6 +189,7 @@ class Tile : public BasicResource {
   std::set<Link*> out_links;
   std::vector<std::unique_ptr<FunctionUnit>> functional_unit_storage; // Owns FUs.
   std::set<FunctionUnit*> functional_units; // Non-owning, for fast lookup.
+  RegisterFileCluster *register_file_cluster = nullptr;
 };
 
 //===----------------------------------------------------------------------===//

include/NeuraDialect/Mapping/HeuristicMapping/HeuristicMapping.h

@@ -13,14 +13,20 @@ class HeuristicMapping : public MappingStrategy {
 public:
   HeuristicMapping(int max_location_to_try = 5, int max_backtrack_depth = 3)
       : max_location_to_try(max_location_to_try), max_backtrack_depth(3) {}
+
   bool map(std::vector<Operation *> &sorted_ops,
            const Architecture &architecture,
            MappingState &mapping_state) override;
+
   std::string getName() const override {
-    if (max_backtrack_depth == 1 && max_location_to_try == INT_MAX) {
+    if (max_location_to_try == 1 &&
+        max_backtrack_depth == 1) {
+      return "simple";
+    } else if (max_location_to_try == INT_MAX &&
+               max_backtrack_depth == 1) {
       return "greedy";
-    } else if (max_backtrack_depth == INT_MAX &&
-               max_location_to_try == INT_MAX) {
+    } else if (max_location_to_try == INT_MAX &&
+               max_backtrack_depth == INT_MAX) {
       return "exhaustive";
     } else {
       return "heuristic";

include/NeuraDialect/Mapping/MappingState.h

@@ -18,13 +18,19 @@ struct MappingLoc {
   int time_step;
 
   bool operator==(const MappingLoc &other) const {
-    return resource == other.resource && time_step == other.time_step;
+    return resource->getKind() == other.resource->getKind() &&
+           resource->getId() == other.resource->getId() &&
+           time_step == other.time_step;
   }
 
   bool operator<(const MappingLoc &other) const {
-    if (time_step != other.time_step)
-      return time_step < other.time_step;
-    return resource->getId() < other.resource->getId();
+    if (resource->getKind() != other.resource->getKind()) {
+      return resource->getKind() < other.resource->getKind();
+    }
+    if (resource->getId() != other.resource->getId()) {
+      return resource->getId() < other.resource->getId();
+    }
+    return time_step < other.time_step;
   }
 };
 
@@ -34,9 +40,10 @@ struct MappingLoc {
 namespace std {
 template <> struct hash<mlir::neura::MappingLoc> {
   std::size_t operator()(const mlir::neura::MappingLoc &loc) const {
-    std::size_t h1 = std::hash<mlir::neura::BasicResource *>()(loc.resource);
-    std::size_t h2 = std::hash<int>()(loc.time_step);
-    return h1 ^ (h2 << 1);
+    std::size_t h1 = std::hash<int>()(static_cast<int>(loc.resource->getKind()));
+    std::size_t h2 = std::hash<int>()(loc.resource->getId());
+    std::size_t h3 = std::hash<int>()(loc.time_step);
+    return h1 ^ (h2 << 1) ^ (h3 << 2);
   }
 };
 } // namespace std
@@ -61,15 +68,21 @@ class MappingState {
   // it will check (tile 2, step 1), (tile 2, step 5), (tile 2, step 9), etc.
   bool isAvailableAcrossTime(const MappingLoc &loc) const;
 
+  // Checks if a hardware resource is available across a time range.
+  // This function leverages the isAvailableAcrossTime function in each
+  // time step.
+  bool isAvailableAcrossTimeInRange(BasicResource *resource,
+                                    int start_time,
+                                    int exclusive_end_time) const;
+
   // Gets the operation at a specific (tile/link, time_step) location.
   std::optional<Operation *> getOpAt(MappingLoc loc) const;
 
+  std::optional<Operation *> getOpAtLocAcrossTime(MappingLoc loc) const;
+
   // Counts the number of operations at a specific resource across time steps.
   int countOpsAtResource(BasicResource *resource) const;
 
-  // Gets all MRRG nodes.
-  const std::set<MappingLoc> &getAllLocs() const;
-
   // Gets all MRRG nodes allocated to a given op.
   const std::vector<MappingLoc> &getAllLocsOfOp(Operation *op) const;
 
@@ -126,7 +139,6 @@ class MappingState {
   int II;
   static constexpr int kMaxSteps = 10;
 
-  std::set<MappingLoc> all_locs;
   std::set<MappingLoc> occupied_locs;
   std::map<MappingLoc, Operation *> loc_to_op;
   std::map<Operation *, std::vector<MappingLoc>> op_to_locs;

include/NeuraDialect/Mapping/mapping_util.h

@@ -84,5 +84,13 @@ bool canReachLocInTime(const std::vector<Operation *> &producers,
                        int deadline_step,
                        const MappingState &mapping_state);
 
+// Gets an available register (for the given time range) in the given tile.
+// The end_time is exclusive, meaning the register should be available
+// until end_time - 1. Returns nullptr if no available register found.
+Register *getAvailableRegister(const MappingState &mapping_state,
+                               Tile *tile,
+                               int start_time,
+                               int exclusive_end_time);
+
 } // namespace neura
 } // namespace mlir

lib/NeuraDialect/Architecture/Architecture.cpp

@@ -42,6 +42,41 @@ const std::set<Link *> &Tile::getOutLinks() const { return out_links; }
 
 const std::set<Link *> &Tile::getInLinks() const { return in_links; }
 
+void Tile::addRegisterFileCluster(RegisterFileCluster* register_file_cluster) {
+  assert(register_file_cluster && "Cannot add null register file cluster");
+  if (this->register_file_cluster != nullptr) {
+    llvm::errs() << "Warning: Overwriting existing register file cluster ("
+                 << this->register_file_cluster->getId() << ") in Tile "
+                 << this->id << "\n";
+  }
+  assert(this->register_file_cluster == nullptr &&
+         "Register file cluster already exists");
+  this->register_file_cluster = register_file_cluster;
+  register_file_cluster->setTile(this);
+}
+
+const RegisterFileCluster* Tile::getRegisterFileCluster() const {
+  return register_file_cluster;
+}
+
+const std::vector<RegisterFile *> Tile::getRegisterFiles() const {
+  std::vector<RegisterFile*> all_register_files;
+  for (const auto& [id, file] : this->register_file_cluster->getRegisterFiles()) {
+    all_register_files.push_back(file);
+  }
+  return all_register_files;
+}
+
+const std::vector<Register *> Tile::getRegisters() const {
+  std::vector<Register *> all_registers;
+  for (const auto& [reg_file_id, reg_file] : this->register_file_cluster->getRegisterFiles()) {
+    for (const auto& [reg_id, reg] : reg_file->getRegisters()) {
+      all_registers.push_back(reg);
+    }
+  }
+  return all_registers;
+}
+
 //===----------------------------------------------------------------------===//
 // Link
 //===----------------------------------------------------------------------===//
@@ -81,18 +116,22 @@ Tile *FunctionUnit::getTile() const {
 // Register
 //===----------------------------------------------------------------------===//
 
-Tile *Register::getTile() const {
-  return this->register_file ? register_file->getTile() : nullptr;
-}
-
 Register::Register(int id) { this->id = id; }
 
 int Register::getId() const { return id; }
 
+Tile *Register::getTile() const {
+  return this->register_file ? register_file->getTile() : nullptr;
+}
+
 void Register::setRegisterFile(RegisterFile* register_file) {
   this->register_file = register_file;
 }
 
+RegisterFile *Register::getRegisterFile() const {
+  return this->register_file;
+}
+
 //===----------------------------------------------------------------------===//
 // Register File
 //===----------------------------------------------------------------------===//
@@ -117,14 +156,15 @@ void RegisterFile::addRegister(Register* reg) {
 const std::map<int, Register*>& RegisterFile::getRegisters() const {
   return this->registers;
 }
-//===----------------------------------------------------------------------===//
-// Register File Cluster
-//===----------------------------------------------------------------------===//
 
 RegisterFileCluster* RegisterFile::getRegisterFileCluster() const {
   return this->register_file_cluster;
 }
 
+//===----------------------------------------------------------------------===//
+// Register File Cluster
+//===----------------------------------------------------------------------===//
+
 RegisterFileCluster::RegisterFileCluster(int id) { this->id = id; }
 
 int RegisterFileCluster::getId() const { return id; }
@@ -153,24 +193,64 @@ const std::map<int, RegisterFile*>& RegisterFileCluster::getRegisterFiles() cons
 Architecture::Architecture(int width, int height) {
   const int num_tiles = width * height;
 
+  // Initializes tiles.
   tile_storage.reserve(num_tiles);
-
-  for (int i = 0; i < width; ++i) {
-    for (int j = 0; j < height; ++j) {
-      const int id = i * width + j;
-      auto tile = std::make_unique<Tile>(id, i, j);
+  for (int y = 0; y < height; ++y) {
+    for (int x = 0; x < width; ++x) {
+      const int id = y * width + x;
+      // const int id = x * width + y;
+      auto tile = std::make_unique<Tile>(id, x, y);
       id_to_tile[id] = tile.get();
-      coord_to_tile[{i, j}] = tile.get();
+      coord_to_tile[{x, y}] = tile.get();
       tile_storage.push_back(std::move(tile));
     }
   }
 
+  // Initializes register file cluster for each tile.
+  int reg_id = 0;
+  for (int y = 0; y < height; ++y) {
+    for (int x = 0; x < width; ++x) {
+      // Gets the tile by coordinates.
+      Tile *tile = getTile(x, y);
+
+      // Creates registers as a register file.
+      // FIXME: We have to assign different IDs due to the hash function
+      // cannot distinguish between different register files..
+      Register *register_0 = new Register(reg_id++);
+      Register *register_1 = new Register(reg_id++);
+      RegisterFile *register_file_0 = new RegisterFile(0);
+      register_file_0->addRegister(register_0);
+      register_file_0->addRegister(register_1);
+
+      // Creates registers as a register file.
+      Register *register_2 = new Register(reg_id++);
+      Register *register_3 = new Register(reg_id++);
+      RegisterFile *register_file_1 = new RegisterFile(1);
+      register_file_1->addRegister(register_2);
+      register_file_1->addRegister(register_3);
+
+      // Assembles register files into a cluster.
+      RegisterFileCluster *register_file_cluster = new RegisterFileCluster(y * width + x);
+      register_file_cluster->addRegisterFile(register_file_0);
+      register_file_cluster->addRegisterFile(register_file_1);
+
+      // Adds register file cluster to the tile.
+      tile->addRegisterFileCluster(register_file_cluster);
+      llvm::errs() << "Tile (" << x << ", " << y
+                 << ") added register file cluster with ID: "
+                 << register_file_cluster->getId() << "\n";
+    }
+  }
+
   // TODO: Model topology based on the architecture specs.
   // https://github.com/coredac/dataflow/issues/52.
   int link_id = 0;
-  for (int i = 0; i < width; ++i) {
-    for (int j = 0; j < height; ++j) {
+  for (int j = 0; j < height; ++j) {
+    for (int i = 0; i < width; ++i) {
+      // Gets the tile by coordinates.
       Tile *tile = getTile(i, j);
+
+      // Creates links to neighboring tiles.
       if (i > 0) {
         auto link_towards_left = std::make_unique<Link>(link_id++);
         link_towards_left->connect(tile, getTile(i - 1, j));