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/SpatialMapping/SpatialMapping.h

@@ -0,0 +1,65 @@
+#ifndef NEURA_SPATIAL_MAPPING_H
+#define NEURA_SPATIAL_MAPPING_H
+
+#include "NeuraDialect/Architecture/Architecture.h"
+#include "NeuraDialect/Mapping/MappingState.h"
+#include "NeuraDialect/Mapping/MappingStrategy.h"
+#include <climits>
+#include <map>
+#include <set>
+#include <vector>
+
+namespace mlir {
+namespace neura {
+// Implements a spatial-only mapping strategy where each tile can only be
+// assigned to one operation regardless of time.
+class SpatialMapping : public MappingStrategy {
+public:
+  SpatialMapping(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 "spatial-simple";
+    } else if (max_location_to_try == INT_MAX && max_backtrack_depth == 1) {
+      return "spatial-greedy";
+    } else if (max_location_to_try == INT_MAX &&
+               max_backtrack_depth == INT_MAX) {
+      return "spatial-exhaustive";
+    } else {
+      return "spatial-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, size_t current_index, int backtrack_depth);
+
+  std::vector<MappingLoc>
+  caculateSpatialAward(Operation *op, std::set<Operation *> &critical_ops,
+                       int target_level, const Architecture &architecture,
+                       const MappingState &mapping_state);
+
+  bool canReachLocSpatial(const std::vector<Operation *> &producers,
+                          const MappingLoc &target_loc, int deadline_step,
+                          const MappingState &mapping_state);
+
+  bool canReachLocSpatial(const MappingLoc &src_loc, const MappingLoc &dst_loc,
+                          int deadline_step, const MappingState &mapping_state);
+
+  // Maximum number of locations to try for each operation.
+  int max_location_to_try;
+  // Maximum depth for backtracking.
+  int max_backtrack_depth;
+};
+} // namespace neura
+} // namespace mlir
+#endif

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,8 @@ add_mlir_dialect_library(MLIRNeura
         Mapping/mapping_util.cpp
         Mapping/MappingState.cpp
         Mapping/MappingStrategy.cpp
-        Mapping/HeuristicMapping/HeuristicMapping.cpp
+        Mapping/SpatialTemporalMapping/SpatialTemporalMapping.cpp
+        Mapping/SpatialMapping/SpatialMapping.cpp
         Architecture/Architecture.cpp
 
         ADDITIONAL_HEADER_DIRS