Refactor the implementation logic of backtrack mapping

include/NeuraDialect/Mapping/HeuristicMapping/HeuristicMapping.h

@@ -1,45 +1,44 @@
-#ifndef NEURA_BACKTRACK_MAPPING_H
-#define NEURA_BACKTRACK_MAPPING_H
+#ifndef NEURA_HEURISTIC_MAPPING_H
+#define NEURA_HEURISTIC_MAPPING_H
 
+#include "NeuraDialect/Mapping/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 HeuristicMapping : public MappingStrategy {
+class HeuristicMapping : public Mapping {
 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) {}
+      : 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) {
+    if (max_location_to_try == 1 && max_backtrack_depth == 1) {
       return "simple";
-    } else if (max_location_to_try == INT_MAX &&
-               max_backtrack_depth == 1) {
+    } 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";
+      return "customized";
     }
   }
 
 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);
+  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
@@ -49,24 +48,4 @@ class HeuristicMapping : public MappingStrategy {
 } // 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
+#endif // NEURA_HEURISTIC_MAPPING_H

include/NeuraDialect/Mapping/Mapping.h

@@ -0,0 +1,29 @@
+#ifndef NEURA_MAPPING_H
+#define NEURA_MAPPING_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 Mapping {
+public:
+  virtual ~Mapping() = default;
+
+  // Applies the mapping strategy to map operations onto hardware
+  virtual bool
+  map(std::vector<std::pair<Operation *, int>> &sorted_ops_with_alap_levels,
+      std::set<Operation *> &critical_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_H

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

lib/NeuraDialect/CMakeLists.txt

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