Enable backtrack mapping strategy

include/NeuraDialect/Mapping/HeuristicMapping/HeuristicMapping.h

@@ -0,0 +1,64 @@
+#ifndef NEURA_BACKTRACK_MAPPING_H
+#define NEURA_BACKTRACK_MAPPING_H
+
+#include "NeuraDialect/Mapping/MappingState.h"
+#include "NeuraDialect/Mapping/MappingStrategy.h"
+#include <climits>
+#include <map>
+#include <set>
+
+namespace mlir {
+namespace neura {
+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) {
+      return "greedy";
+    } else if (max_backtrack_depth == INT_MAX &&
+               max_location_to_try == INT_MAX) {
+      return "exhaustive";
+    } else {
+      return "heuristic";
+    }
+  }
+
+private:
+  bool mapWithBacktrack(std::vector<Operation *> &sorted_ops,
+                        const Architecture &architecture,
+                        MappingState &mapping_state, size_t current_index,
+                        int backtrack_depth);
+
+  // Configuration parameters.
+  int max_location_to_try; // Maximum number of locations to try for
+                           // each op
+  int max_backtrack_depth; // Maximum depth for backtracking
+};
+} // namespace neura
+} // namespace mlir
+
+namespace mlir {
+namespace neura {
+class MappingStateSnapshot {
+public:
+  MappingStateSnapshot(const MappingState &mapping_state);
+
+  void restore(MappingState &mapping_state);
+
+  std::map<Operation *, std::vector<MappingLoc>> getOpToLocs() {
+    return this->op_to_locs;
+  }
+
+private:
+  std::set<MappingLoc> occupied_locs;
+  std::map<MappingLoc, Operation *> loc_to_op;
+  std::map<Operation *, std::vector<MappingLoc>> op_to_locs;
+};
+} // namespace neura
+} // namespace mlir
+
+#endif // NEURA_BACKTRACK_MAPPING_H

include/NeuraDialect/Mapping/MappingState.h

@@ -1,20 +1,20 @@
 #ifndef NEURA_MAPPING_STATE_H
 #define NEURA_MAPPING_STATE_H
 
-#include "mlir/IR/Operation.h"
 #include "NeuraDialect/Architecture/Architecture.h"
+#include "mlir/IR/Operation.h"
 #include "llvm/Support/raw_ostream.h"
+#include <optional>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
-#include <optional>
 
 namespace mlir {
 namespace neura {
 
 // Represents a spatial-temporal location: (resource, time_step)
 struct MappingLoc {
-  BasicResource* resource;
+  BasicResource *resource;
   int time_step;
 
   bool operator==(const MappingLoc &other) const {
@@ -32,15 +32,14 @@ struct MappingLoc {
 } // namespace mlir
 
 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);
+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);
   }
 };
-}
+} // namespace std
 
 namespace mlir {
 namespace neura {
@@ -63,7 +62,7 @@ class MappingState {
   bool isAvailableAcrossTime(const MappingLoc &loc) const;
 
   // Gets the operation at a specific (tile/link, time_step) location.
-  std::optional<Operation*> getOpAt(MappingLoc loc) const;
+  std::optional<Operation *> getOpAt(MappingLoc loc) const;
 
   // Counts the number of operations at a specific resource across time steps.
   int countOpsAtResource(BasicResource *resource) const;
@@ -83,11 +82,11 @@ class MappingState {
   // Gets neighboring tiles on next step of a given MappingLoc.
   std::vector<MappingLoc> getNextStepTiles(MappingLoc loc) const;
 
-//   // Gets neighboring links on next step of a given MappingLoc.
-//   const std::vector<MappingLoc> &getNextStepLinks(MappingLoc loc) const;
+  //   // Gets neighboring links on next step of a given MappingLoc.
+  //   const std::vector<MappingLoc> &getNextStepLinks(MappingLoc loc) const;
 
-//   // Gets neighboring tiles on current step of a given MappingLoc.
-//   const std::vector<MappingLoc> &getCurrentStepTiles(MappingLoc loc) const;
+  //   // Gets neighboring tiles on current step of a given MappingLoc.
+  //   const std::vector<MappingLoc> &getCurrentStepTiles(MappingLoc loc) const;
 
   // Gets neighboring links on current step of a given MappingLoc.
   std::vector<MappingLoc> getCurrentStepLinks(MappingLoc loc) const;
@@ -100,15 +99,37 @@ class MappingState {
 
   void dumpOpToLocs(llvm::raw_ostream &os = llvm::errs()) const;
 
+  // Getters for state information.
+  const std::set<MappingLoc> &getOccupiedLocs() const {
+    return this->occupied_locs;
+  }
+  const std::map<MappingLoc, Operation *> &getLocToOp() const {
+    return this->loc_to_op;
+  }
+  const std::map<Operation *, std::vector<MappingLoc>> &getOpToLocs() const {
+    return this->op_to_locs;
+  }
+
+  // Setters for state information.
+  void setOccupiedLocs(const std::set<MappingLoc> &locs) {
+    this->occupied_locs = locs;
+  }
+  void setLocToOp(const std::map<MappingLoc, Operation *> &loc_to_op) {
+    this->loc_to_op = loc_to_op;
+  }
+  void setOpToLocs(const std::map<Operation *, std::vector<MappingLoc>> &op_to_locs) {
+    this->op_to_locs = op_to_locs;
+  }
+
 private:
   // Initiation interval.
   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;
+  std::map<MappingLoc, Operation *> loc_to_op;
+  std::map<Operation *, std::vector<MappingLoc>> op_to_locs;
 };
 
 } // namespace neura

include/NeuraDialect/Mapping/MappingStrategy.h

@@ -0,0 +1,28 @@
+#ifndef NEURA_MAPPING_STRATEGY_H
+#define NEURA_MAPPING_STRATEGY_H
+
+#include "NeuraDialect/Architecture/Architecture.h"
+#include "NeuraDialect/Mapping/MappingState.h"
+#include <vector>
+
+namespace mlir {
+namespace neura {
+
+// Abstract base class for different mapping strategies.
+class MappingStrategy {
+public:
+  virtual ~MappingStrategy() = default;
+
+  // Applies the mapping strategy to map operations onto hardware
+  virtual bool map(std::vector<Operation *> &sorted_ops,
+                   const Architecture &architecture,
+                   MappingState &mapping_state) = 0;
+
+  // Gets the name of this strategy
+  virtual std::string getName() const = 0;
+};
+
+} // namespace neura
+} // namespace mlir
+
+#endif // NEURA_MAPPING_STRATEGY_H

lib/NeuraDialect/Architecture/Architecture.cpp

@@ -12,59 +12,37 @@ Tile::Tile(int id, int x, int y) {
   this->y = y;
 }
 
-int Tile::getId() const {
-  return id;
-}
+int Tile::getId() const { return id; }
 
-int Tile::getX() const {
-  return x;
-}
+int Tile::getX() const { return x; }
 
-int Tile::getY() const {
-  return y;
-}
+int Tile::getY() const { return y; }
 
-void Tile::linkDstTile(Link* link, Tile* tile) {
+void Tile::linkDstTile(Link *link, Tile *tile) {
   assert(tile && "Cannot link to a null tile");
   dst_tiles.insert(tile);
   out_links.insert(link);
   tile->src_tiles.insert(this);
   tile->in_links.insert(link);
 }
 
-const std::set<Tile*>& Tile::getDstTiles() const {
-  return dst_tiles;
-}
+const std::set<Tile *> &Tile::getDstTiles() const { return dst_tiles; }
 
-const std::set<Tile*>& Tile::getSrcTiles() const {
-  return src_tiles;
-}
+const std::set<Tile *> &Tile::getSrcTiles() const { return src_tiles; }
 
-const std::set<Link*>& Tile::getOutLinks() const {
-  return out_links;
-}
+const std::set<Link *> &Tile::getOutLinks() const { return out_links; }
 
-const std::set<Link*>& Tile::getInLinks() const {
-  return in_links;
-}
+const std::set<Link *> &Tile::getInLinks() const { return in_links; }
 
-Link::Link(int id) {
-  this->id = id;
-}
+Link::Link(int id) { this->id = id; }
 
-int Link::getId() const {
-  return id;
-}
+int Link::getId() const { return id; }
 
-Tile* Link::getSrcTile() const {
-  return src_tile;
-}
+Tile *Link::getSrcTile() const { return src_tile; }
 
-Tile* Link::getDstTile() const {
-  return dst_tile;
-}
+Tile *Link::getDstTile() const { return dst_tile; }
 
-void Link::connect(Tile* src, Tile* dst) {
+void Link::connect(Tile *src, Tile *dst) {
   assert(src && dst && "Cannot connect null tiles");
   src_tile = src;
   dst_tile = dst;
@@ -91,7 +69,7 @@ Architecture::Architecture(int width, int height) {
   int link_id = 0;
   for (int i = 0; i < width; ++i) {
     for (int j = 0; j < height; ++j) {
-      Tile* tile = getTile(i, j);
+      Tile *tile = getTile(i, j);
       if (i > 0) {
         auto link_towards_left = std::make_unique<Link>(link_id++);
         link_towards_left->connect(tile, getTile(i - 1, j));
@@ -101,7 +79,7 @@ Architecture::Architecture(int width, int height) {
         auto link_towards_right = std::make_unique<Link>(link_id++);
         link_towards_right->connect(tile, getTile(i + 1, j));
         link_storage.push_back(std::move(link_towards_right));
-    }
+      }
       if (j > 0) {
         auto link_towards_down = std::make_unique<Link>(link_id++);
         link_towards_down->connect(tile, getTile(i, j - 1));
@@ -116,20 +94,20 @@ Architecture::Architecture(int width, int height) {
   }
 }
 
-Tile* Architecture::getTile(int id) {
+Tile *Architecture::getTile(int id) {
   auto it = id_to_tile.find(id);
   assert(it != id_to_tile.end() && "Tile with given ID not found");
   return it->second;
 }
 
-Tile* Architecture::getTile(int x, int y) {
+Tile *Architecture::getTile(int x, int y) {
   auto it = coord_to_tile.find({x, y});
   assert(it != coord_to_tile.end() && "Tile with given coordinates not found");
   return it->second;
 }
 
-std::vector<Tile*> Architecture::getAllTiles() const {
-  std::vector<Tile*> result;
+std::vector<Tile *> Architecture::getAllTiles() const {
+  std::vector<Tile *> result;
   for (auto &tile : tile_storage)
     result.push_back(tile.get());
   return result;
@@ -139,8 +117,8 @@ int Architecture::getNumTiles() const {
   return static_cast<int>(id_to_tile.size());
 }
 
-std::vector<Link*> Architecture::getAllLinks() const {
-  std::vector<Link*> all_links;
+std::vector<Link *> Architecture::getAllLinks() const {
+  std::vector<Link *> all_links;
   for (const auto &link : link_storage) {
     all_links.push_back(link.get());
   }

lib/NeuraDialect/CMakeLists.txt

@@ -5,6 +5,8 @@ add_mlir_dialect_library(MLIRNeura
         NeuraPasses.cpp
         Mapping/mapping_util.cpp
         Mapping/MappingState.cpp
+        Mapping/MappingStrategy.cpp
+        Mapping/HeuristicMapping/HeuristicMapping.cpp
         Architecture/Architecture.cpp
 
         ADDITIONAL_HEADER_DIRS