Refactor the implementation logic of backtrack mapping

include/NeuraDialect/Mapping/MappingState.h

@@ -40,7 +40,8 @@ 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<int>()(static_cast<int>(loc.resource->getKind()));
+    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);
@@ -71,8 +72,7 @@ class MappingState {
   // 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,
+  bool isAvailableAcrossTimeInRange(BasicResource *resource, int start_time,
                                     int exclusive_end_time) const;
 
   // Gets the operation at a specific (tile/link, time_step) location.
@@ -130,7 +130,8 @@ class MappingState {
   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) {
+  void setOpToLocs(
+      const std::map<Operation *, std::vector<MappingLoc>> &op_to_locs) {
     this->op_to_locs = op_to_locs;
   }
 
@@ -147,4 +148,24 @@ class MappingState {
 } // 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_MAPPING_STATE_H

include/NeuraDialect/Mapping/SpatialTemporalMapping/SpatialTemporalMapping.h

@@ -0,0 +1,52 @@
+#ifndef NEURA_SPATIALTEMPORAL_MAPPING_H
+#define NEURA_SPATIALTEMPORAL_MAPPING_H
+
+#include "NeuraDialect/Mapping/MappingState.h"
+#include "NeuraDialect/Mapping/MappingStrategy.h"
+#include "NeuraDialect/Mapping/mapping_util.h"
+#include <climits>
+#include <map>
+#include <set>
+
+namespace mlir {
+namespace neura {
+class SpatialTemporalMapping : public MappingStrategy {
+public:
+  SpatialTemporalMapping(int max_location_to_try = 5,
+                         int max_backtrack_depth = 3)
+      : max_location_to_try(max_location_to_try),
+        max_backtrack_depth(max_backtrack_depth) {}
+
+  bool map(std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
+           std::set<Operation *> &critical_ops,
+           const Architecture &architecture,
+           MappingState &mapping_state) override;
+
+  std::string getName() const override {
+    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_location_to_try == INT_MAX &&
+               max_backtrack_depth == INT_MAX) {
+      return "exhaustive";
+    } else {
+      return "heuristic";
+    }
+  }
+
+private:
+  bool mapWithBacktrack(
+      std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
+      std::set<Operation *> &critical_ops, const Architecture &architecture,
+      MappingState &mapping_state);
+
+  // 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
+
+#endif // NEURA_SPATIALTEMPORAL_MAPPING_H

lib/NeuraDialect/CMakeLists.txt

@@ -6,7 +6,7 @@ add_mlir_dialect_library(MLIRNeura
         Mapping/mapping_util.cpp
         Mapping/MappingState.cpp
         Mapping/MappingStrategy.cpp
-        Mapping/HeuristicMapping/HeuristicMapping.cpp
+        Mapping/SpatialTemporalMapping/SpatialTemporalMapping.cpp
         Architecture/Architecture.cpp
 
         ADDITIONAL_HEADER_DIRS

lib/NeuraDialect/Mapping/MappingState.cpp

@@ -63,7 +63,8 @@ std::optional<Operation *> MappingState::getOpAt(MappingLoc loc) const {
   return it->second;
 }
 
-std::optional<Operation *> MappingState::getOpAtLocAcrossTime(MappingLoc loc) const {
+std::optional<Operation *>
+MappingState::getOpAtLocAcrossTime(MappingLoc loc) const {
   for (int t = loc.time_step % II; t < II * kMaxSteps; t += II) {
     MappingLoc check_loc = {loc.resource, t};
     auto it = loc_to_op.find(check_loc);
@@ -84,7 +85,8 @@ int MappingState::countOpsAtResource(BasicResource *resource) const {
   return count;
 }
 
-const std::vector<MappingLoc> &MappingState::getAllLocsOfOp(Operation *op) const {
+const std::vector<MappingLoc> &
+MappingState::getAllLocsOfOp(Operation *op) const {
   auto it = op_to_locs.find(op);
   if (it != op_to_locs.end()) {
     return it->second;
@@ -130,7 +132,8 @@ std::vector<MappingLoc> MappingState::getNextStepTiles(MappingLoc loc) const {
 //   return it != current_step_tiles.end() ? it->second : empty;
 // }
 
-std::vector<MappingLoc> MappingState::getCurrentStepLinks(MappingLoc loc) const {
+std::vector<MappingLoc>
+MappingState::getCurrentStepLinks(MappingLoc loc) const {
   assert((loc.resource->getKind() == ResourceKind::Tile) &&
          "Current step links can only be queried for tiles");
   std::vector<MappingLoc> current_step_links;
@@ -276,4 +279,16 @@ void MappingState::encodeMappingState() {
     }
     op->setAttr("mapping_locs", mlir::ArrayAttr::get(ctx, mapping_entries));
   }
-}
+}
+
+MappingStateSnapshot::MappingStateSnapshot(const MappingState &mapping_state) {
+  this->occupied_locs = mapping_state.getOccupiedLocs();
+  this->loc_to_op = mapping_state.getLocToOp();
+  this->op_to_locs = mapping_state.getOpToLocs();
+}
+
+void MappingStateSnapshot::restore(MappingState &mapping_state) {
+  mapping_state.setOccupiedLocs(this->occupied_locs);
+  mapping_state.setLocToOp(this->loc_to_op);
+  mapping_state.setOpToLocs(this->op_to_locs);
+}