[Support] Add relative path caching to InstancePathCache (#8489)

This commit adds infrastructure for caching relative instance paths
between modules in the InstanceGraph. It adds a relativePathsCache field
to InstancePathCache to store paths between arbitrary module pairs

Author: uenoku

include/circt/Support/InstanceGraph.h

@@ -345,17 +345,21 @@ inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
   return os;
 }
 
-/// A data structure that caches and provides absolute paths to module
+/// A data structure that caches and provides paths to module
 /// instances in the IR.
 struct InstancePathCache {
   /// The instance graph of the IR.
   InstanceGraph &instanceGraph;
 
   explicit InstancePathCache(InstanceGraph &instanceGraph)
       : instanceGraph(instanceGraph) {}
+
+  // Return all absolute paths from the top-level node to the given module.
   ArrayRef<InstancePath> getAbsolutePaths(ModuleOpInterface op);
-  ArrayRef<InstancePath> getAbsolutePaths(ModuleOpInterface op,
-                                          InstanceGraphNode *top);
+
+  // Return all relative paths from the given node to the given module.
+  ArrayRef<InstancePath> getRelativePaths(ModuleOpInterface op,
+                                          InstanceGraphNode *node);
 
   /// Replace an InstanceOp. This is required to keep the cache updated.
   void replaceInstance(InstanceOpInterface oldOp, InstanceOpInterface newOp);
@@ -367,11 +371,18 @@ struct InstancePathCache {
   InstancePath prependInstance(InstanceOpInterface inst, InstancePath path);
 
 private:
+  using PathsCache = DenseMap<Operation *, ArrayRef<InstancePath>>;
+  ArrayRef<InstancePath> getPaths(ModuleOpInterface op, InstanceGraphNode *top,
+                                  PathsCache &cache);
+
   /// An allocator for individual instance paths and entire path lists.
   llvm::BumpPtrAllocator allocator;
 
   /// Cached absolute instance paths.
-  DenseMap<Operation *, ArrayRef<InstancePath>> absolutePathsCache;
+  PathsCache absolutePathsCache;
+
+  /// Cached relative instance paths.
+  DenseMap<InstanceGraphNode *, PathsCache> relativePathsCache;
 };
 
 } // namespace igraph

lib/Support/InstanceGraph.cpp

@@ -226,36 +226,46 @@ static InstancePath empty{};
 
 ArrayRef<InstancePath>
 InstancePathCache::getAbsolutePaths(ModuleOpInterface op) {
-  return getAbsolutePaths(op, instanceGraph.getTopLevelNode());
+  return getPaths(op, instanceGraph.getTopLevelNode(), absolutePathsCache);
 }
 
-// NOLINTBEGIN(misc-no-recursion)
 ArrayRef<InstancePath>
-InstancePathCache::getAbsolutePaths(ModuleOpInterface op,
-                                    InstanceGraphNode *top) {
+InstancePathCache::getRelativePaths(ModuleOpInterface op,
+                                    InstanceGraphNode *node) {
+  if (node == instanceGraph.getTopLevelNode())
+    return getAbsolutePaths(op);
+  return getPaths(op, node, relativePathsCache[node]);
+}
+
+// NOLINTBEGIN(misc-no-recursion)
+ArrayRef<InstancePath> InstancePathCache::getPaths(ModuleOpInterface op,
+                                                   InstanceGraphNode *top,
+                                                   PathsCache &cache) {
   InstanceGraphNode *node = instanceGraph[op];
 
   if (node == top) {
     return empty;
   }
 
   // Fast path: hit the cache.
-  auto cached = absolutePathsCache.find(op);
-  if (cached != absolutePathsCache.end())
+  auto cached = cache.find(op);
+  if (cached != cache.end())
     return cached->second;
 
   // For each instance, collect the instance paths to its parent and append the
   // instance itself to each.
   SmallVector<InstancePath, 8> extendedPaths;
   for (auto *inst : node->uses()) {
     if (auto module = inst->getParent()->getModule()) {
-      auto instPaths = getAbsolutePaths(module);
+      auto instPaths = getPaths(module, top, cache);
       extendedPaths.reserve(instPaths.size());
       for (auto path : instPaths) {
         extendedPaths.push_back(appendInstance(
             path, cast<InstanceOpInterface>(*inst->getInstance())));
       }
-    } else {
+    } else if (inst->getParent() == top) {
+      // Special case when `inst` is a top-level instance and
+      // `inst->getParent()` is a pseudo top-level node.
       extendedPaths.emplace_back(empty);
     }
   }
@@ -267,7 +277,7 @@ InstancePathCache::getAbsolutePaths(ModuleOpInterface op,
     std::copy(extendedPaths.begin(), extendedPaths.end(), paths);
     pathList = ArrayRef<InstancePath>(paths, extendedPaths.size());
   }
-  absolutePathsCache.insert({op, pathList});
+  cache.insert({op, pathList});
   return pathList;
 }
 // NOLINTEND(misc-no-recursion)
@@ -331,29 +341,33 @@ void InstancePathCache::replaceInstance(InstanceOpInterface oldOp,
     return llvm::any_of(
         paths, [&](InstancePath p) { return llvm::is_contained(p, oldOp); });
   };
-
-  for (auto &iter : absolutePathsCache) {
-    if (!instanceExists(iter.getSecond()))
-      continue;
-    SmallVector<InstancePath, 8> updatedPaths;
-    for (auto path : iter.getSecond()) {
-      const auto *iter = llvm::find(path, oldOp);
-      if (iter == path.end()) {
-        // path does not contain the oldOp, just copy it as is.
-        updatedPaths.push_back(path);
+  auto updateCache = [&](PathsCache &cache) {
+    for (auto &iter : cache) {
+      if (!instanceExists(iter.getSecond()))
         continue;
+      SmallVector<InstancePath, 8> updatedPaths;
+      for (auto path : iter.getSecond()) {
+        const auto *iter = llvm::find(path, oldOp);
+        if (iter == path.end()) {
+          // path does not contain the oldOp, just copy it as is.
+          updatedPaths.push_back(path);
+          continue;
+        }
+        auto *newPath = allocator.Allocate<InstanceOpInterface>(path.size());
+        llvm::copy(path, newPath);
+        newPath[iter - path.begin()] = newOp;
+        updatedPaths.push_back(InstancePath(ArrayRef(newPath, path.size())));
       }
-      auto *newPath = allocator.Allocate<InstanceOpInterface>(path.size());
-      llvm::copy(path, newPath);
-      newPath[iter - path.begin()] = newOp;
-      updatedPaths.push_back(InstancePath(ArrayRef(newPath, path.size())));
+      // Move the list of paths into the bump allocator for later quick
+      // retrieval.
+      auto *paths = allocator.Allocate<InstancePath>(updatedPaths.size());
+      llvm::copy(updatedPaths, paths);
+      iter.getSecond() = ArrayRef<InstancePath>(paths, updatedPaths.size());
     }
-    // Move the list of paths into the bump allocator for later quick
-    // retrieval.
-    auto *paths = allocator.Allocate<InstancePath>(updatedPaths.size());
-    llvm::copy(updatedPaths, paths);
-    iter.getSecond() = ArrayRef<InstancePath>(paths, updatedPaths.size());
-  }
+  };
+  updateCache(absolutePathsCache);
+  for (auto &iter : relativePathsCache)
+    updateCache(iter.getSecond());
 }
 
 #include "circt/Support/InstanceGraphInterface.cpp.inc"

unittests/Dialect/HW/InstancePathTest.cpp

@@ -40,6 +40,23 @@ TEST(InstancePathTest, Enumerate) {
   ASSERT_EQ(1ull, topPaths.size());
 }
 
+TEST(InstancePathTest, RelativePath) {
+  MLIRContext context;
+  ModuleOp circuit = fixtures::createModule(&context);
+  hw::InstanceGraph graph(circuit);
+  igraph::InstancePathCache pathCache(graph);
+
+  auto cat = cast<HWModuleOp>(*circuit.getBody()->rbegin());
+  auto alligator = cast<HWModuleOp>(*std::next(circuit.getBody()->begin()));
+  auto catToAlligatorPaths = pathCache.getRelativePaths(cat, graph[alligator]);
+
+  ASSERT_EQ(1ull, catToAlligatorPaths.size());
+
+  ASSERT_EQ(2ull, catToAlligatorPaths[0].size());
+  EXPECT_EQ("bear", catToAlligatorPaths[0][0].getInstanceName());
+  EXPECT_EQ("cat", catToAlligatorPaths[0][1].getInstanceName());
+}
+
 TEST(InstancePathTest, AppendPrependInstance) {
   MLIRContext context;
   ModuleOp circuit = fixtures::createModule(&context);