Refactor the implementation logic of backtrack mapping

include/NeuraDialect/Mapping/SpatialTemporalMapping/SpatialTemporalMapping.h → include/NeuraDialect/Mapping/BacktrackMapping/BacktrackMapping.h

@@ -1,19 +1,18 @@
-#ifndef NEURA_SPATIALTEMPORAL_MAPPING_H
-#define NEURA_SPATIALTEMPORAL_MAPPING_H
+#ifndef NEURA_BACKTRACK_MAPPING_H
+#define NEURA_BACKTRACK_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 SpatialTemporalMapping : public MappingStrategy {
+class BacktrackMapping : public Mapping {
 public:
-  SpatialTemporalMapping(int max_location_to_try = 5,
-                         int max_backtrack_depth = 3)
+  BacktrackMapping(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) {}
 
@@ -49,4 +48,4 @@ class SpatialTemporalMapping : public MappingStrategy {
 } // namespace neura
 } // namespace mlir
 
-#endif // NEURA_SPATIALTEMPORAL_MAPPING_H
+#endif // NEURA_BACKTRACK_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

lib/NeuraDialect/CMakeLists.txt

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

lib/NeuraDialect/Mapping/SpatialTemporalMapping/SpatialTemporalMapping.cpp → lib/NeuraDialect/Mapping/BacktrackMapping/BacktrackMapping.cpp

@@ -1,4 +1,4 @@
-#include "NeuraDialect/Mapping/SpatialTemporalMapping/SpatialTemporalMapping.h"
+#include "NeuraDialect/Mapping/BacktrackMapping/BacktrackMapping.h"
 #include "NeuraDialect/Mapping/MappingState.h"
 #include "NeuraDialect/Mapping/mapping_util.h"
 #include "NeuraDialect/NeuraOps.h"
@@ -13,7 +13,7 @@ inline bool is_non_materialized(Operation *op) {
   return mlir::isa<neura::ReserveOp, neura::CtrlMovOp, neura::DataMovOp>(op);
 }
 
-bool SpatialTemporalMapping::map(
+bool BacktrackMapping::map(
     std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
     std::set<Operation *> &critical_ops, const Architecture &architecture,
     MappingState &mapping_state) {
@@ -22,12 +22,12 @@ bool SpatialTemporalMapping::map(
                           mapping_state);
 }
 
-bool SpatialTemporalMapping::mapWithBacktrack(
+bool BacktrackMapping::mapWithBacktrack(
     std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
     std::set<Operation *> &critical_ops, const Architecture &architecture,
     MappingState &mapping_state) {
   llvm::outs() << "---------------------------------------------------------\n";
-  llvm::outs() << "[SpatialTemporalMapping] Starting mapping with "
+  llvm::outs() << "[BacktrackMapping] Starting mapping with "
                << sorted_ops_with_levels.size() << " operations.\n";
   llvm::outs() << "Configuration: MAX Backtrack Depth = "
                << this->max_backtrack_depth

lib/NeuraDialect/Mapping/Mapping.cpp

@@ -0,0 +1,2 @@
+#include "NeuraDialect/Mapping/Mapping.h"
+#include <memory>

lib/NeuraDialect/Transforms/MapToAcceleratorPass.cpp

@@ -1,8 +1,9 @@
 #include <deque>
+#include <memory>
 
 #include "NeuraDialect/Architecture/Architecture.h"
+#include "NeuraDialect/Mapping/BacktrackMapping/BacktrackMapping.h"
 #include "NeuraDialect/Mapping/MappingState.h"
-#include "NeuraDialect/Mapping/SpatialTemporalMapping/SpatialTemporalMapping.h"
 #include "NeuraDialect/Mapping/mapping_util.h"
 #include "NeuraDialect/NeuraDialect.h"
 #include "NeuraDialect/NeuraOps.h"
@@ -12,6 +13,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
@@ -41,63 +43,69 @@ struct MapToAcceleratorPass
   Option<std::string> mappingStrategy{
       *this, "mapping-strategy",
       llvm::cl::desc("Mapping strategy to use for mapping operations to the "
-                     "accelerator. Options: greedy, exhaustive, "
-                     "heuristic=max_loc,max_depth (default "
-                     "max_loc=5, max_depth=3)"),
+                     "accelerator. Options: backtrack (default)."),
+      llvm::cl::init("backtrack")};
+  Option<std::string> backtrackConfig{
+      *this, "backtrack-config",
+      llvm::cl::desc(
+          "Backtrack configuration used for mapping operations to the "
+          "accelerator. Options: simple, greedy, exhaustive, "
+          "heuristic=max_loc,max_depth (default "
+          "max_loc=5, max_depth=3)"),
       llvm::cl::init("heuristic")};
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
-
+    std::unique_ptr<Mapping> mapping;
     StringRef mappingStrategy_stringRef(mappingStrategy.getValue());
-    // Creates a mapping strategy based on the provided option.
-    std::unique_ptr<MappingStrategy> mapping_strategy;
-    if (mappingStrategy_stringRef == "simple") {
-      mapping_strategy = std::make_unique<SpatialTemporalMapping>(1, 1);
-    } else if (mappingStrategy_stringRef == "greedy") {
-      mapping_strategy = std::make_unique<SpatialTemporalMapping>(INT_MAX, 1);
-    } else if (mappingStrategy_stringRef == "exhaustive") {
-      mapping_strategy =
-          std::make_unique<SpatialTemporalMapping>(INT_MAX, INT_MAX);
-    } else if (mappingStrategy_stringRef == "heuristic") {
-      mapping_strategy = std::make_unique<SpatialTemporalMapping>(
-          5, 3); // Randomly picked default values for max_loc and max_depth
-    } else if (mappingStrategy_stringRef.starts_with("heuristic=")) {
-      // Used for custom backtrack parameters.
-      // Example: "heuristic=5,3" means max_loc=5, max_depth=3
-      // Extracts the parameters after "heuristic=".
-      StringRef paramsRef =
-          mappingStrategy_stringRef.substr(strlen("heuristic="));
-      size_t comma_pos = paramsRef.find(',');
-
-      if (comma_pos != StringRef::npos) {
-        StringRef max_loc_str = paramsRef.substr(0, comma_pos);
-        StringRef max_depth_str = paramsRef.substr(comma_pos + 1);
-
-        int max_loc, max_depth;
-        if (!max_loc_str.getAsInteger(10, max_loc) &&
-            !max_depth_str.getAsInteger(10, max_depth)) {
-          mapping_strategy =
-              std::make_unique<SpatialTemporalMapping>(max_loc, max_depth);
-          llvm::errs()
-              << "[MapToAcceleratorPass] Use custom backtrack parameters: "
-              << "max_location_to_try=" << max_loc
-              << ", max_backtrack_depth=" << max_depth << "\n";
+    StringRef backtrackConfig_stringRef(backtrackConfig.getValue());
+    if (mappingStrategy_stringRef == "backtrack" ||
+        mappingStrategy_stringRef.empty()) {
+
+      if (backtrackConfig_stringRef == "simple") {
+        mapping = std::make_unique<BacktrackMapping>(1, 1);
+      } else if (backtrackConfig_stringRef == "greedy") {
+        mapping = std::make_unique<BacktrackMapping>(INT_MAX, 1);
+      } else if (backtrackConfig_stringRef == "exhaustive") {
+        mapping = std::make_unique<BacktrackMapping>(INT_MAX, INT_MAX);
+      } else if (backtrackConfig_stringRef == "heuristic") {
+        mapping = std::make_unique<BacktrackMapping>(5, 3);
+      } else if (backtrackConfig_stringRef.starts_with("heuristic=")) {
+        // Used for custom backtrack parameters.
+        // Example: "heuristic=5,3" means max_loc=5, max_depth=3
+        // Extracts the parameters after "heuristic=".
+        StringRef paramsRef =
+            backtrackConfig_stringRef.substr(strlen("heuristic="));
+        size_t comma_pos = paramsRef.find(',');
+
+        if (comma_pos != StringRef::npos) {
+          StringRef max_loc_str = paramsRef.substr(0, comma_pos);
+          StringRef max_depth_str = paramsRef.substr(comma_pos + 1);
+
+          int max_loc, max_depth;
+          if (!max_loc_str.getAsInteger(10, max_loc) &&
+              !max_depth_str.getAsInteger(10, max_depth)) {
+            mapping = std::make_unique<BacktrackMapping>(max_loc, max_depth);
+            llvm::errs()
+                << "[MapToAcceleratorPass] Use custom backtrack parameters: "
+                << "max_location_to_try=" << max_loc
+                << ", max_backtrack_depth=" << max_depth << "\n";
+          } else {
+            llvm::errs() << "[MapToAcceleratorPass] Illegal backtrack "
+                            "parameters format: "
+                         << backtrackConfig_stringRef << "\n";
+            return;
+          }
         } else {
           llvm::errs()
               << "[MapToAcceleratorPass] Illegal backtrack parameters format: "
-              << mappingStrategy << "\n";
+              << backtrackConfig_stringRef << "\n";
           return;
         }
-      } else {
-        llvm::errs()
-            << "[MapToAcceleratorPass] Illegal backtrack parameters format: "
-            << mappingStrategy << "\n";
-        return;
       }
     } else {
       llvm::errs() << "[MapToAcceleratorPass] Unsupported mapping strategy: "
-                   << mappingStrategy << "\n";
+                   << mappingStrategy_stringRef << "\n";
       return;
     }
 
@@ -184,8 +192,8 @@ struct MapToAcceleratorPass
             << "\n";
         // Creates a mapping state for the current II.
         MappingState mapping_state(architecture, ii);
-        if (mapping_strategy->map(sorted_ops_with_alap_levels, critical_ops,
-                                  architecture, mapping_state)) {
+        if (mapping->map(sorted_ops_with_alap_levels, critical_ops,
+                         architecture, mapping_state)) {
           // success
           llvm::errs() << "[MapToAcceleratorPass] Successfully mapped function "
                        << func.getName() << "' with II = " << ii << "\n";

test/controflow_fuse/simple_loop/simple_loop.mlir

@@ -58,7 +58,7 @@
 // RUN: --fuse-control-flow \
 // RUN: --fold-constant \
 // RUN: --insert-data-mov \
-// RUN: --map-to-accelerator="mapping-strategy=heuristic" | FileCheck %s -check-prefix=FUSE-MAPPING
+// RUN: --map-to-accelerator="mapping-strategy=backtrack backtrack-config=heuristic" | FileCheck %s -check-prefix=FUSE-MAPPING
 
 module attributes {} {
   func.func @_Z11simple_loopPiS_(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {

test/controflow_fuse/simple_loop_reduction/simple_loop_reduction.mlir

@@ -57,7 +57,7 @@
 // RUN: --fuse-control-flow \
 // RUN: --fold-constant \
 // RUN: --insert-data-mov \
-// RUN: --map-to-accelerator="mapping-strategy=heuristic=4,5" | FileCheck %s -check-prefix=FUSE-MAPPING
+// RUN: --map-to-accelerator="mapping-strategy=backtrack backtrack-config=heuristic=4,5" | FileCheck %s -check-prefix=FUSE-MAPPING
 
 module attributes {} {
   func.func @_Z10simpleloopv() -> i32 attributes {llvm.linkage = #llvm.linkage<external>} {

test/mapping_quality/branch_for.mlir

@@ -46,7 +46,7 @@
 // RUN:   --transform-ctrl-to-data-flow \
 // RUN:   --fold-constant \
 // RUN:   --insert-data-mov \
-// RUN:   --map-to-accelerator="mapping-strategy=heuristic" \
+// RUN:   --map-to-accelerator="mapping-strategy=backtrack backtrack-config=heuristic" \
 // RUN:   | FileCheck %s -check-prefix=MAPPING
 
 // RUN: mlir-neura-opt %s \

test/neura/ctrl/branch_for.mlir

@@ -46,7 +46,7 @@
 // RUN:   --transform-ctrl-to-data-flow \
 // RUN:   --fold-constant \
 // RUN:   --insert-data-mov \
-// RUN:   --map-to-accelerator="mapping-strategy=simple" \
+// RUN:   --map-to-accelerator="mapping-strategy=backtrack backtrack-config=simple" \
 // RUN:   | FileCheck %s -check-prefix=MAPPING
 
 // RUN: mlir-neura-opt %s \
@@ -57,7 +57,7 @@
 // RUN:   --transform-ctrl-to-data-flow \
 // RUN:   --fold-constant \
 // RUN:   --insert-data-mov \
-// RUN:   --map-to-accelerator="mapping-strategy=simple" \
+// RUN:   --map-to-accelerator="mapping-strategy=backtrack backtrack-config=simple" \
 // RUN:   --generate-code
 // RUN: FileCheck %s --input-file=generated-instructions.json -check-prefix=INST