Register Bypass Constraint

include/NeuraDialect/Mapping/MappingState.h

@@ -5,17 +5,18 @@
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>
+#include <unordered_map>
 #include <vector>
 
 namespace mlir {
 namespace neura {
 
 // Occupy status for multi-cycle pipeline support.
 // These states define how a tile/FU is occupied at a given time step.
-#define SINGLE_OCCUPY     0 // A single-cycle op is in the FU (exclusive)
+#define SINGLE_OCCUPY 0     // A single-cycle op is in the FU (exclusive)
 #define START_PIPE_OCCUPY 1 // A multi-cycle op starts in the FU
-#define END_PIPE_OCCUPY   2 // A multi-cycle op ends in the FU
-#define IN_PIPE_OCCUPY    3 // A multi-cycle op is occupying the FU (pipelined)
+#define END_PIPE_OCCUPY 2   // A multi-cycle op ends in the FU
+#define IN_PIPE_OCCUPY 3    // A multi-cycle op is occupying the FU (pipelined)
 
 // Represents a spatial-temporal location: (resource, time_step)
 struct MappingLoc {
@@ -83,16 +84,17 @@ class MappingState {
   // Note that the check is performed in II granularity.
   // For example, if II is 4, and we want to check (tile 2, step 5), then
   // it will check (tile 2, step 1), (tile 2, step 5), (tile 2, step 9), etc.
-  bool isAvailableAcrossTime(const MappingLoc &loc) const;
+  bool isAvailableAcrossTime(const MappingLoc &loc,
+                             Operation *op = nullptr) const;
 
   // Checks if a location is available for a specific occupy status.
   // This implements the pipeline-aware availability checking:
   // - SINGLE_OCCUPY: only available if location is completely free
   // - START_PIPE_OCCUPY: available if free or IN_PIPE_OCCUPY or END_PIPE_OCCUPY
   // - END_PIPE_OCCUPY: available if free or IN_PIPE_OCCUPY or START_PIPE_OCCUPY
   // - IN_PIPE_OCCUPY: always available (can pipeline with any status)
-  bool isAvailableForOccupyStatus(const MappingLoc &loc,
-                                  int new_occupy_status) const;
+  bool isAvailableForOccupyStatus(const MappingLoc &loc, int new_occupy_status,
+                                  Operation *op = nullptr) const;
 
   // Gets the occupy status at a specific location across time domain.
   // Returns -1 if the location is not occupied.
@@ -102,7 +104,20 @@ class MappingState {
   // This function leverages the isAvailableAcrossTime function in each
   // time step.
   bool isAvailableAcrossTimeInRange(BasicResource *resource, int start_time,
-                                    int exclusive_end_time) const;
+                                    int exclusive_end_time,
+                                    Operation *op = nullptr) const;
+
+  // Checks availability of a register's cluster write port across the relevant
+  // time steps.  Returns false if a DIFFERENT register in the same
+  // RegisterFile is already writing at a congruent time slot.  Multiple writes
+  // to the SAME register are allowed (idempotent).
+  bool isRegisterWriteAvailableAcrossTime(Register *reg, int time_step) const;
+
+  // Checks availability of a register's cluster read port across the relevant
+  // time steps.  Returns false if a DIFFERENT register in the same
+  // RegisterFile is already reading at a congruent time slot.  Multiple reads
+  // from the SAME register are allowed (shared read).
+  bool isRegisterReadAvailableAcrossTime(Register *reg, int time_step) const;
 
   // Gets the operation at a specific (tile/link, time_step) location.
   std::optional<Operation *> getOpAt(MappingLoc loc) const;
@@ -152,6 +167,17 @@ class MappingState {
   const std::map<Operation *, std::vector<MappingLoc>> &getOpToLocs() const {
     return this->op_to_locs;
   }
+  const std::unordered_map<
+      RegisterFile *, std::unordered_map<int, std::pair<Register *, int>>> &
+  getRegFileWriteToOccupyOperations() const {
+    return this->reg_file_write_to_occupy_operations;
+  }
+
+  const std::unordered_map<
+      RegisterFile *, std::unordered_map<int, std::pair<Register *, int>>> &
+  getRegFileReadToOccupyOperations() const {
+    return this->reg_file_read_to_occupy_operations;
+  }
 
   // Setters for state information.
   void setOccupiedLocs(
@@ -166,6 +192,19 @@ class MappingState {
       const std::map<Operation *, std::vector<MappingLoc>> &op_to_locs) {
     this->op_to_locs = op_to_locs;
   }
+  void setRegFileWriteToOccupyOperations(
+      const std::unordered_map<
+          RegisterFile *, std::unordered_map<int, std::pair<Register *, int>>>
+          &records) {
+    this->reg_file_write_to_occupy_operations = records;
+  }
+
+  void setRegFileReadToOccupyOperations(
+      const std::unordered_map<
+          RegisterFile *, std::unordered_map<int, std::pair<Register *, int>>>
+          &records) {
+    this->reg_file_read_to_occupy_operations = records;
+  }
 
 private:
   // Initiation interval.
@@ -178,6 +217,42 @@ class MappingState {
   std::map<MappingLoc, std::vector<std::pair<int, Operation *>>> occupied_locs;
   std::map<MappingLoc, Operation *> loc_to_op;
   std::map<Operation *, std::vector<MappingLoc>> op_to_locs;
+
+  // Record table for register cluster write occupancy, keyed by (RegisterFile*,
+  // time_step % II). Maps to the Register* that is writing to the cluster at
+  // that canonical time slot. At most one writer is allowed per cluster per
+  // slot, enforcing the hardware constraint that a register cluster supports
+  // only a single write port per cycle.  However, multiple writes to the SAME
+  // register are allowed (idempotent). Storing the reference count as the pair
+  // value ensures proper tracking during overlaps or backtracking.
+  std::unordered_map<RegisterFile *,
+                     std::unordered_map<int, std::pair<Register *, int>>>
+      reg_file_write_to_occupy_operations;
+
+  // Record table for register cluster read occupancy, keyed by (RegisterFile*,
+  // time_step % II). Maps to the Register* that is reading from the cluster at
+  // that canonical time slot. Multiple reads from the SAME register are allowed
+  // (shared read), but reads from DIFFERENT registers in the same cluster are
+  // rejected (only one read port).
+  std::unordered_map<RegisterFile *,
+                     std::unordered_map<int, std::pair<Register *, int>>>
+      reg_file_read_to_occupy_operations;
+
+  // Records the write operation for a register cluster slot.
+  // Returns false if the cluster slot is already occupied by a writer.
+  bool addWriteToRegFileRecord(Register *reg, int time_step);
+
+  // Records the read operation for a register cluster slot.
+  // Returns false if the cluster slot is already occupied by a reader.
+  bool addReadToRegFileRecord(Register *reg, int time_step);
+
+  // Removes the write operation from a register cluster slot
+  // when the op is unbound/released.
+  void removeWriteFromRegFileRecord(Register *reg, int time_step);
+
+  // Removes the read operation from a register cluster slot
+  // when the op is unbound/released.
+  void removeReadFromRegFileRecord(Register *reg, int time_step);
 };
 
 } // namespace neura
@@ -199,6 +274,12 @@ class MappingStateSnapshot {
   std::map<MappingLoc, std::vector<std::pair<int, Operation *>>> occupied_locs;
   std::map<MappingLoc, Operation *> loc_to_op;
   std::map<Operation *, std::vector<MappingLoc>> op_to_locs;
+  std::unordered_map<RegisterFile *,
+                     std::unordered_map<int, std::pair<Register *, int>>>
+      reg_file_write_to_occupy_operations;
+  std::unordered_map<RegisterFile *,
+                     std::unordered_map<int, std::pair<Register *, int>>>
+      reg_file_read_to_occupy_operations;
 };
 } // namespace neura
 } // namespace mlir

include/NeuraDialect/Mapping/mapping_util.h

@@ -113,8 +113,18 @@ bool canReachLocInTime(const std::vector<Operation *> &producers,
 // Gets an available register (for the given time range) in the given tile.
 // The end_time is exclusive, meaning the register should be available
 // until end_time - 1. Returns nullptr if no available register found.
+//
+// The optional `op` parameter is the DataMovOp being routed. It is forwarded
+// to MappingState::isAvailableAcrossTime() / isAvailableForOccupyStatus() so
+// that multiple DataMovOps carrying the same materialized source value can
+// share a single physical register. When `op` is non-null the availability
+// check recognises that two DataMovOps reading the identical value do not
+// actually conflict, because the single register read port broadcasts the
+// value to all consumers. Passing nullptr disables this sharing and falls
+// back to the strict one-occupant-per-register rule.
 Register *getAvailableRegister(const MappingState &mapping_state, Tile *tile,
-                               int start_time, int exclusive_end_time);
+                               int start_time, int exclusive_end_time,
+                               Operation *op = nullptr);
 
 // Gets the execution latency of an operation from its "latency" attribute.
 // Returns 1 (single-cycle) if the attribute is not present.