Add heuristic for no-producer operations

include/NeuraDialect/Mapping/HeuristicMapping/HeuristicMapping.h

@@ -16,9 +16,7 @@ class HeuristicMapping : public Mapping {
       : 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,
+  bool map(const Architecture &architecture,
            MappingState &mapping_state) override;
 
   std::string getName() const override {
@@ -34,18 +32,57 @@ class HeuristicMapping : public Mapping {
     }
   }
 
+  // Temporary structure to hold the result of no-producer operation mapping.
+  struct NoProducerOpCandidate {
+    MappingStateSnapshot state;
+    int mapped_ops_count;
+    bool fully_mapped;
+    int current_ii;
+
+    // Calculates the quality score of the solution.
+    double getQualityScore() const {
+      // If it is not fully mapped, return the number of mapped operations.
+      if (!fully_mapped) {
+        return mapped_ops_count;
+      }
+
+      // If it is fully mapped, return a score inversely proportional to II.
+      return 1000.0 - this->current_ii * 100.0;
+    }
+  };
+
 private:
-  bool mapWithBacktrack(
-      std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
-      std::set<Operation *> &critical_ops, const Architecture &architecture,
-      MappingState &mapping_state);
-
-  // Gets the sorted candidate locations for a given operation based on spatial
-  // execution model.
-  std::vector<MappingLoc>
-  calculateSpatialAward(Operation *op, std::set<Operation *> &critical_ops,
-                        int target_level, const Architecture &architecture,
-                        const MappingState &mapping_state);
+  bool mapWithBacktrack(const Architecture &architecture,
+                        MappingState &mapping_state);
+
+  // Checks if the current operation is after a no-producer operation.
+  bool
+  isAfterNoProducerOp(const std::unordered_set<Operation *> &no_producer_ops,
+                      int current_op_index);
+
+  // Performs backtracking to restore the previous mapping state.
+  bool performBacktrack(const std::unordered_set<Operation *> &no_producer_ops,
+                        std::vector<MappingStateSnapshot> &snapshots,
+                        std::vector<int> &candidate_history,
+                        std::vector<int> &operation_index_history,
+                        int current_op_index, MappingState &mapping_state);
+
+  // Attempts to map a no-producer operation with global exploration.
+  NoProducerOpCandidate
+  tryToMapNoProducerOp(Operation *current_op, int current_op_index,
+                       const std::vector<MappingLoc> &candidate_locs,
+                       std::vector<MappingStateSnapshot> &snapshots,
+                       std::vector<int> &candidate_history,
+                       std::vector<int> &operation_index_history,
+                       const Architecture &architecture,
+                       MappingState &mapping_state);
+
+  // Evaluates a candidate location for a no-producer operation.
+  NoProducerOpCandidate evaluateNoProducerOpCandidate(
+      Operation *current_op, int current_op_index,
+      const MappingLoc &candidate_loc, int candidate_index,
+      int total_candidates, const Architecture &architecture,
+      MappingState &mapping_state, MappingStateSnapshot &initial_state);
 
   // Configuration parameters.
   // Maximum number of locations to try for each op.

include/NeuraDialect/Mapping/Mapping.h

@@ -14,15 +14,28 @@ class Mapping {
   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;
+  virtual bool map(const Architecture &architecture,
+                   MappingState &mapping_state) = 0;
 
   // Gets the name of this strategy
   virtual std::string getName() const = 0;
 
+  void loadDfg(
+      const std::vector<std::pair<Operation *, int>> &sorted_ops_with_levels,
+      const std::set<Operation *> &critical_ops);
+  std::vector<std::pair<Operation *, int>> getSortedOpsWithLevels() const {
+    return this->sorted_ops_with_levels;
+  }
+  std::vector<std::pair<Operation *, int>>
+  getMaterializedOpsWithLevels() const {
+    return this->materialized_ops_with_levels;
+  }
+  std::set<Operation *> getCriticalOps() const { return this->critical_ops; }
+
 private:
+  std::vector<std::pair<Operation *, int>> sorted_ops_with_levels;
+  std::vector<std::pair<Operation *, int>> materialized_ops_with_levels;
+  std::set<Operation *> critical_ops;
 };
 
 } // namespace neura

include/NeuraDialect/Mapping/mapping_util.h

@@ -79,11 +79,10 @@ bool placeAndRoute(Operation *op, const MappingLoc &target_loc,
 
 // Calculates the award of mapping locations for a given op, the returned
 // locations are sorted based on the award.
-std::vector<MappingLoc> calculateAward(Operation *op,
-                                       std::set<Operation *> &critical_ops,
-                                       int target_level,
-                                       const Architecture &architecture,
-                                       const MappingState &mapping_state);
+std::vector<MappingLoc>
+calculateAward(Operation *op, const std::set<Operation *> &critical_ops,
+               int target_level, const Architecture &architecture,
+               const MappingState &mapping_state);
 
 void updateAward(std::map<MappingLoc, int> &locs_with_award, MappingLoc loc,
                  int award);