Merge pull request #56 from coredac/mrrg

Construct MRRG

Author: tancheng

include/NeuraDialect/Architecture/Architecture.h

@@ -77,6 +77,12 @@ class Link : public BasicResource {
   Tile* dst_tile;
 };
 
+struct PairHash {
+  std::size_t operator()(const std::pair<int, int> &coord) const {
+    return std::hash<int>()(coord.first) ^ (std::hash<int>()(coord.second) << 1);
+  }
+};
+
 /// Describes the entire CGRA architecture.
 class Architecture {
 public:
@@ -87,12 +93,16 @@ class Architecture {
 
   Link* getLink(int id);
 
-  int getNumTiles() const { return static_cast<int>(tiles.size()); }
-  std::vector<Tile*> getAllTiles() { return tiles; }
+  int getNumTiles() const;
+  std::vector<Tile*> getAllTiles() const;
+  std::vector<Link*> getAllLinks() const;
 
 private:
-  std::vector<std::unique_ptr<Tile>> tileStorage;
-  std::vector<Tile*> tiles;
+  std::vector<std::unique_ptr<Tile>> tile_storage;
+//   std::vector<Tile*> tiles;
+  std::vector<std::unique_ptr<Link>> link_storage;
+  std::unordered_map<int, Tile*> id_to_tile;
+  std::unordered_map<std::pair<int, int>, Tile*, PairHash> coord_to_tile;
 };
 
 } // namespace neura

include/NeuraDialect/Mapping/MappingState.h

@@ -11,22 +11,74 @@
 namespace mlir {
 namespace neura {
 
-// Represents a spatial-temporal location: (resource, timeStep)
-using MappingLoc = std::pair<BasicResource*, int>;
+// Represents a spatial-temporal location: (resource, time_step)
+struct MappingLoc {
+  BasicResource* resource;
+  int time_step;
+
+  bool operator==(const MappingLoc &other) const {
+    return resource == other.resource && time_step == other.time_step;
+  }
+};
+
+} // namespace neura
+} // namespace mlir
+
+namespace std {
+template <>
+struct hash<mlir::neura::MappingLoc> {
+  std::size_t operator()(const mlir::neura::MappingLoc& loc) const {
+    std::size_t h1 = std::hash<mlir::neura::BasicResource*>()(loc.resource);
+    std::size_t h2 = std::hash<int>()(loc.time_step);
+    return h1 ^ (h2 << 1);
+  }
+};
+}
+
+namespace mlir {
+namespace neura {
 
 // Tracks placement and routing of ops on the CGRA.
 class MappingState {
 public:
-  // Binds a (tile/link, timeStep) location to an operation.
+  MappingState(const Architecture &arch, int II);
+  // Binds a (tile/link, time_step) location to an operation.
   void bindOp(MappingLoc loc, Operation *op);
 
-  // Checks if a (tile/link, timeStep) is available (unoccupied).
+  // Checks if a (tile/link, time_step) is available (unoccupied).
   bool isAvailable(const MappingLoc &loc) const;
 
-  // Gets the operation at a specific (tile/link, timeStep) location.
+  // Gets the operation at a specific (tile/link, time_step) location.
   std::optional<Operation*> getOpAt(MappingLoc loc) const;
 
+  // Gets all MRRG nodes.
+  const std::unordered_set<MappingLoc> &getAllLocs() const;
+
+  // Gets neighboring tiles on next step of a given MappingLoc.
+  const std::vector<MappingLoc> &getNextStepTiles(MappingLoc loc) const;
+
+  // Gets neighboring links on next step of a given MappingLoc.
+  const std::vector<MappingLoc> &getNextStepLinks(MappingLoc loc) const;
+
+  // Gets neighboring tiles on current step of a given MappingLoc.
+  const std::vector<MappingLoc> &getCurrentStepTiles(MappingLoc loc) const;
+
+  // Gets neighboring links on current step of a given MappingLoc.
+  const std::vector<MappingLoc> &getCurrentStepLinks(MappingLoc loc) const;
+
+  int getII() const { return II; }
+
 private:
+  // Initiation interval.
+  int II;
+  std::unordered_set<MappingLoc> all_locs;
+  // current and next step tiles and links for a given MappingLoc. Note that
+  // the key MappingLoc is either a pair of (tile, time_step) or (link, time_step).
+  std::unordered_map<MappingLoc, std::vector<MappingLoc>> next_step_tiles;
+  std::unordered_map<MappingLoc, std::vector<MappingLoc>> next_step_links;
+  std::unordered_map<MappingLoc, std::vector<MappingLoc>> current_step_tiles;
+  std::unordered_map<MappingLoc, std::vector<MappingLoc>> current_step_links;
+
   std::unordered_map<MappingLoc, Operation*> loc_to_op;
   std::unordered_set<MappingLoc> occupied_locs;
 };

lib/NeuraDialect/Architecture/Architecture.cpp

@@ -1,4 +1,6 @@
 #include "NeuraDialect/Architecture/Architecture.h"
+#include "llvm/Support/raw_ostream.h"
+
 #include <cassert>
 
 using namespace mlir;
@@ -72,14 +74,15 @@ void Link::connect(Tile* src, Tile* dst) {
 Architecture::Architecture(int width, int height) {
   const int num_tiles = width * height;
 
-  tileStorage.reserve(num_tiles);
-  tiles.reserve(num_tiles);
+  tile_storage.reserve(num_tiles);
 
   for (int i = 0; i < width; ++i) {
     for (int j = 0; j < height; ++j) {
-      auto tile = std::make_unique<Tile>(i * width + j, i, j);
-      tiles.push_back(tile.get());
-      tileStorage.push_back(std::move(tile));
+      const int id = i * width + j;
+      auto tile = std::make_unique<Tile>(id, i, j);
+      id_to_tile[id] = tile.get();
+      coord_to_tile[{i, j}] = tile.get();
+      tile_storage.push_back(std::move(tile));
     }
   }
 
@@ -92,39 +95,54 @@ Architecture::Architecture(int width, int height) {
       if (i > 0) {
         auto link_towards_left = std::make_unique<Link>(link_id++);
         link_towards_left->connect(tile, getTile(i - 1, j));
+        link_storage.push_back(std::move(link_towards_left));
       }
       if (i < width - 1) {
         auto link_towards_right = std::make_unique<Link>(link_id++);
         link_towards_right->connect(tile, getTile(i + 1, j));
+        link_storage.push_back(std::move(link_towards_right));
     }
       if (j > 0) {
         auto link_towards_down = std::make_unique<Link>(link_id++);
         link_towards_down->connect(tile, getTile(i, j - 1));
+        link_storage.push_back(std::move(link_towards_down));
       }
       if (j < height - 1) {
         auto link_towards_up = std::make_unique<Link>(link_id++);
         link_towards_up->connect(tile, getTile(i, j + 1));
+        link_storage.push_back(std::move(link_towards_up));
       }
     }
   }
 }
 
 Tile* Architecture::getTile(int id) {
-  for (const auto &tile : tiles) {
-    if (tile->getId() == id) {
-      return tile;
-    }
-  }
-  assert(false && "Tile with given ID not found");
-  return nullptr;
+  auto it = id_to_tile.find(id);
+  assert(it != id_to_tile.end() && "Tile with given ID not found");
+  return it->second;
 }
 
 Tile* Architecture::getTile(int x, int y) {
-  for (const auto &tile : tiles) {
-    if (tile->getX() == x && tile->getY() == y) {
-      return tile;
-    }
+  auto it = coord_to_tile.find({x, y});
+  assert(it != coord_to_tile.end() && "Tile with given coordinates not found");
+  return it->second;
+}
+
+std::vector<Tile*> Architecture::getAllTiles() const {
+  std::vector<Tile*> result;
+  for (auto &tile : tile_storage)
+    result.push_back(tile.get());
+  return result;
+}
+
+int Architecture::getNumTiles() const {
+  return static_cast<int>(id_to_tile.size());
+}
+
+std::vector<Link*> Architecture::getAllLinks() const {
+  std::vector<Link*> all_links;
+  for (const auto &link : link_storage) {
+    all_links.push_back(link.get());
   }
-  assert(false && "Tile with given coordinates not found");
-  return nullptr;
+  return all_links;
 }

lib/NeuraDialect/CMakeLists.txt

@@ -20,6 +20,7 @@ add_mlir_dialect_library(MLIRNeura
         NeuraTypes.cpp
         NeuraPasses.cpp
         Mapping/mapping_util.cpp
+        Mapping/MappingState.cpp
         Architecture/Architecture.cpp
 
         ADDITIONAL_HEADER_DIRS

lib/NeuraDialect/Mapping/MappingState.cpp

@@ -3,11 +3,67 @@
 using namespace mlir;
 using namespace mlir::neura;
 
+MappingState::MappingState(const Architecture &arch, int II) : II(II) {
+  for (Tile* tile : arch.getAllTiles()) {
+    for (int t = 0; t < II; ++t) {
+      MappingLoc loc = {tile, t};
+      all_locs.insert(loc);
+
+      // Create edges to neighboring tiles at t+1.
+      for (Tile* dst : tile->getDstTiles()) {
+        MappingLoc next_step_dst_tile_loc = {dst, (t + 1) % II}; // modulo II for reuse
+        next_step_tiles[loc].push_back(next_step_dst_tile_loc);
+      }
+
+      // TODO: Not sure whether we need the link on t or t+1.
+      // Creates edges to neighboring links at t.
+      for (Link* dst : tile->getOutLinks()) {
+        MappingLoc current_step_dst_link_loc = {dst, t % II};
+        next_step_tiles[loc].push_back(current_step_dst_link_loc);
+      }
+    }
+  }
+}
+
 void MappingState::bindOp(MappingLoc loc, Operation *op) {
   loc_to_op[loc] = op;
   occupied_locs.insert(loc);
 }
 
 bool MappingState::isAvailable(const MappingLoc &loc) const {
-  return !occupied_locs.contains(loc);
+  return occupied_locs.find(loc) == occupied_locs.end();
+}
+
+std::optional<Operation*> MappingState::getOpAt(MappingLoc loc) const {
+  auto it = loc_to_op.find(loc);
+  if (it == loc_to_op.end()) return std::nullopt;
+  return it->second;
+}
+
+const std::unordered_set<MappingLoc> &MappingState::getAllLocs() const {
+  return all_locs;
 }
+
+const std::vector<MappingLoc> &MappingState::getNextStepTiles(MappingLoc loc) const {
+  static const std::vector<MappingLoc> empty;
+  auto it = next_step_tiles.find(loc);
+  return it != next_step_tiles.end() ? it->second : empty;
+}
+
+const std::vector<MappingLoc> &MappingState::getNextStepLinks(MappingLoc loc) const {
+  static const std::vector<MappingLoc> empty;
+  auto it = next_step_links.find(loc);
+  return it != next_step_links.end() ? it->second : empty;
+}
+
+const std::vector<MappingLoc> &MappingState::getCurrentStepTiles(MappingLoc loc) const {
+  static const std::vector<MappingLoc> empty;
+  auto it = current_step_tiles.find(loc);
+  return it != current_step_tiles.end() ? it->second : empty;
+}
+
+const std::vector<MappingLoc> &MappingState::getCurrentStepLinks(MappingLoc loc) const {
+  static const std::vector<MappingLoc> empty;
+  auto it = current_step_links.find(loc);
+  return it != current_step_links.end() ? it->second : empty;
+}

lib/NeuraDialect/Transforms/MapToAcceleratorPass.cpp

@@ -5,6 +5,7 @@
 #include "NeuraDialect/NeuraOps.h"
 #include "NeuraDialect/NeuraTypes.h"
 #include "NeuraDialect/NeuraPasses.h"
+#include "NeuraDialect/Mapping/MappingState.h"
 #include "NeuraDialect/Mapping/mapping_util.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/PatternMatch.h"
@@ -44,6 +45,7 @@ struct MapToAcceleratorPass
       // Collects and reports recurrence cycles found in the function.
       auto recurrence_cycles = collectRecurrenceCycles(func);
       RecurrenceCycle *longest = nullptr;
+      int rec_mii = 1;
       for (auto &cycle : recurrence_cycles) {
         if (!longest || cycle.length > longest->length)
           longest = &cycle;
@@ -52,10 +54,12 @@ struct MapToAcceleratorPass
       if (longest) {
         llvm::errs() << "[MapToAcceleratorPass] Longest recurrence cycle (length "
                     << longest->length << "):\n";
-        for (Operation *op : longest->operations)
+        for (Operation *op : longest->operations) {
           op->print(llvm::errs()), llvm::errs() << "\n";
+        }
+        rec_mii = longest->length;
         IntegerAttr rec_mii_attr = IntegerAttr::get(
-            IntegerType::get(func.getContext(), 32), longest->length);
+            IntegerType::get(func.getContext(), 32), rec_mii);
         func->setAttr("RecMII", rec_mii_attr);
       }
 
@@ -66,6 +70,11 @@ struct MapToAcceleratorPass
           IntegerType::get(func.getContext(), 32), res_mii);
       func->setAttr("ResMII", res_mii_attr);
 
+      const int minII = std::min(rec_mii, res_mii);
+      constexpr int maxII = 5;
+      for (int ii = minII; ii <= maxII; ++ii) {
+        MappingState state(architecture, ii);
+      }
     });
   }
 };