simplify away cache with holes

Author: zyx-billy

mlir/lib/Target/LLVMIR/DebugImporter.cpp

@@ -319,139 +319,8 @@ getRecSelfConstructor(llvm::DINode *node) {
       .Default(CtorType());
 }
 
-/// An iterative attribute replacer that also handles pruning recursive types.
-/// - Recurses down the attribute tree while replacing attributes based on the
-///   provided replacement map.
-/// - Keeps track of the currently open recursive declarations, and upon
-///   encountering a duplicate declaration, replace with a self-reference.
-static Attribute replaceAndPruneRecursiveTypes(
-    Attribute baseNode,
-    const DenseMap<DIRecursiveTypeAttrInterface, DINodeAttr> &mapping) {
-
-  struct ReplacementState {
-    // The attribute being replaced.
-    Attribute node;
-    // The nested elements of `node`.
-    SmallVector<Attribute> attrs;
-    SmallVector<Type> types;
-    // The current attr being walked on (index into `attrs`).
-    size_t attrIndex;
-    // Whether or not any attr was replaced.
-    bool changed;
-
-    void replaceCurrAttr(Attribute attr) {
-      Attribute &target = attrs[attrIndex];
-      if (attr != target) {
-        changed = true;
-        target = attr;
-      }
-    }
-  };
-
-  // Every iteration, perform replacement on the attribute at `attrIndex` at the
-  // top of `workStack`.
-  // If `attrIndex` reaches past the size of `attrs`, replace `node` with
-  // `attrs` & `types`, and pop off a stack frame and assign the replacement to
-  // the attr being replaced on the previous stack frame.
-  SmallVector<ReplacementState> workStack;
-  workStack.push_back({nullptr, {baseNode}, {}, 0, false});
-
-  // Replacement cache that remembers the context in which the cache is valid.
-  // All unboundRecIds must be in openDecls at time of lookup.
-  struct CacheWithContext {
-    DenseSet<DistinctAttr> unboundRecIds;
-    Attribute entry;
-  };
-  DenseMap<Attribute, CacheWithContext> replacementCache;
-
-  DenseSet<DistinctAttr> openDecls;
-  while (workStack.size() > 1 || workStack.back().attrIndex == 0) {
-    ReplacementState &state = workStack.back();
-
-    // Check for popping condition.
-    if (state.attrIndex == state.attrs.size()) {
-      Attribute result = state.node;
-      if (state.changed)
-        result = result.replaceImmediateSubElements(state.attrs, state.types);
-
-      // Reset context.
-      if (auto recType = dyn_cast<DIRecursiveTypeAttrInterface>(state.node))
-        if (DistinctAttr recId = recType.getRecId())
-          openDecls.erase(recId);
-
-      replacementCache[state.node] = CacheWithContext{openDecls, result};
-
-      workStack.pop_back();
-      ReplacementState &prevState = workStack.back();
-      prevState.replaceCurrAttr(result);
-      ++prevState.attrIndex;
-      continue;
-    }
-
-    Attribute node = state.attrs[state.attrIndex];
-
-    // Lookup in cache first.
-    if (auto it = replacementCache.find(node); it != replacementCache.end()) {
-      // If all the requried recIds are open decls, use cache.
-      if (llvm::set_is_subset(it->second.unboundRecIds, openDecls)) {
-        state.replaceCurrAttr(it->second.entry);
-        ++state.attrIndex;
-        continue;
-      }
-
-      // Otherwise, the cache entry is stale and can be removed now.
-      replacementCache.erase(it);
-    }
-
-    if (auto recType = dyn_cast<DIRecursiveTypeAttrInterface>(node)) {
-      if (DistinctAttr recId = recType.getRecId()) {
-        if (recType.isRecSelf()) {
-          // Replace selfRef based on the provided mapping and re-walk from the
-          // replacement node (do not increment attrIndex).
-          if (DINodeAttr replacement = mapping.lookup(recType)) {
-            state.replaceCurrAttr(replacement);
-            continue;
-          }
-
-          // Otherwise, nothing to do. Advance to next attr.
-          ++state.attrIndex;
-          continue;
-        }
-
-        // Configure context.
-        auto [_, inserted] = openDecls.insert(recId);
-        if (!inserted) {
-          // This is a nested decl. Replace with recSelf. Nothing more to do.
-          state.replaceCurrAttr(recType.getRecSelf(recId));
-          ++state.attrIndex;
-          continue;
-        }
-      }
-    }
-
-    // Recurse into this node.
-    workStack.push_back({node, {}, {}, 0, false});
-    ReplacementState &newState = workStack.back();
-    node.walkImmediateSubElements(
-        [&newState](Attribute attr) { newState.attrs.push_back(attr); },
-        [&newState](Type type) { newState.types.push_back(type); });
-  };
-
-  return workStack.back().attrs.front();
-}
-
 DINodeAttr DebugImporter::RecursionPruner::pruneOrPushTranslationStack(
     llvm::DINode *node) {
-  // Lookup the cache first.
-  auto [result, unboundSelfRefs] = lookup(node);
-  if (result) {
-    // Need to inject unbound self-refs into the previous layer.
-    if (!unboundSelfRefs.empty())
-      translationStack.back().second.unboundSelfRefs.insert(
-          unboundSelfRefs.begin(), unboundSelfRefs.end());
-    return result;
-  }
-
   // If the node type is capable of being recursive, check if it's seen
   // before.
   auto recSelfCtor = getRecSelfConstructor(node);
@@ -465,7 +334,11 @@ DINodeAttr DebugImporter::RecursionPruner::pruneOrPushTranslationStack(
       // a different self-reference. Use that if possible.
       DIRecursiveTypeAttrInterface recSelf = iter->second.recSelf;
       if (!recSelf) {
-        DistinctAttr recId = DistinctAttr::create(UnitAttr::get(context));
+        DistinctAttr recId = nodeToRecId.lookup(node);
+        if (!recId) {
+          recId = DistinctAttr::create(UnitAttr::get(context));
+          nodeToRecId[node] = recId;
+        }
         recSelf = recSelfCtor(recId);
         iter->second.recSelf = recSelf;
       }
@@ -474,7 +347,8 @@ DINodeAttr DebugImporter::RecursionPruner::pruneOrPushTranslationStack(
       return cast<DINodeAttr>(recSelf);
     }
   }
-  return nullptr;
+
+  return lookup(node);
 }
 
 std::pair<DINodeAttr, bool>
@@ -493,39 +367,18 @@ DebugImporter::RecursionPruner::finalizeTranslation(llvm::DINode *node,
   if (DIRecursiveTypeAttrInterface recSelf = state.recSelf) {
     auto recType = cast<DIRecursiveTypeAttrInterface>(result);
     result = cast<DINodeAttr>(recType.withRecId(recSelf.getRecId()));
-
     // Remove this recSelf from the set of unbound selfRefs.
     state.unboundSelfRefs.erase(recSelf);
-
-    // Insert the newly resolved recursive type into the cache entries that
-    // rely on it.
-    // Only need to look at the caches at this level.
-    uint64_t numRemaining = state.cacheSize;
-    for (CachedTranslation &cacheEntry :
-         llvm::make_second_range(llvm::reverse(cache))) {
-      if (numRemaining == 0)
-        break;
-      --numRemaining;
-
-      if (auto refIter = cacheEntry.pendingReplacements.find(recSelf);
-          refIter != cacheEntry.pendingReplacements.end())
-        refIter->second = result;
-    }
   }
 
-  // Insert the current result into the cache.
-  state.cacheSize++;
-  auto [iter, inserted] = cache.try_emplace(node);
-  assert(inserted && "invalid state: caching the same DINode twice");
-  iter->second.attr = result;
-
-  // If this node had any unbound self-refs free when it is registered into
-  // the cache, set up replacement placeholders: This result will need these
-  // unbound self-refs to be replaced before being used.
-  for (DIRecursiveTypeAttrInterface selfRef : state.unboundSelfRefs)
-    iter->second.pendingReplacements.try_emplace(selfRef, nullptr);
-
-  return {result, state.unboundSelfRefs.empty()};
+  // Insert the result into our internal cache if it's not self-contained.
+  if (!state.unboundSelfRefs.empty()) {
+    auto [_, inserted] = dependentCache.try_emplace(
+        node, DependentTranslation{result, state.unboundSelfRefs});
+    assert(inserted && "invalid state: caching the same DINode twice");
+    return {result, false};
+  }
+  return {result, true};
 }
 
 void DebugImporter::RecursionPruner::popTranslationStack(llvm::DINode *node) {
@@ -540,56 +393,30 @@ void DebugImporter::RecursionPruner::popTranslationStack(llvm::DINode *node) {
   if (translationStack.size() == 1) {
     assert(currLayerState.unboundSelfRefs.empty() &&
            "internal error: unbound recursive self reference at top level.");
-    assert(currLayerState.cacheSize == cache.size() &&
-           "internal error: inconsistent cache size");
     translationStack.pop_back();
-    cache.clear();
     return;
   }
 
   // Copy unboundSelfRefs down to the previous level.
   TranslationState &nextLayerState = (++translationStack.rbegin())->second;
   nextLayerState.unboundSelfRefs.insert(currLayerState.unboundSelfRefs.begin(),
                                         currLayerState.unboundSelfRefs.end());
-
-  // The current layer cache is now considered part of the lower layer cache.
-  nextLayerState.cacheSize += currLayerState.cacheSize;
-
-  // Finally pop off this layer when all bookkeeping is done.
   translationStack.pop_back();
 }
 
-std::pair<DINodeAttr, DenseSet<DIRecursiveTypeAttrInterface>>
-DebugImporter::RecursionPruner::lookup(llvm::DINode *node) {
-  auto cacheIter = cache.find(node);
-  if (cacheIter == cache.end())
+DINodeAttr DebugImporter::RecursionPruner::lookup(llvm::DINode *node) {
+  auto cacheIter = dependentCache.find(node);
+  if (cacheIter == dependentCache.end())
     return {};
 
-  CachedTranslation &entry = cacheIter->second;
-
-  if (entry.pendingReplacements.empty())
-    return std::make_pair(entry.attr, DenseSet<DIRecursiveTypeAttrInterface>{});
-
-  Attribute replacedAttr =
-      replaceAndPruneRecursiveTypes(entry.attr, entry.pendingReplacements);
-  DINodeAttr result = cast<DINodeAttr>(replacedAttr);
-
-  // Update cache entry to save replaced version and remove already-applied
-  // replacements.
-  entry.attr = result;
-  DenseSet<DIRecursiveTypeAttrInterface> unboundRefs;
-  DenseSet<DIRecursiveTypeAttrInterface> boundRefs;
-  for (auto [refSelf, replacement] : entry.pendingReplacements) {
-    if (replacement)
-      boundRefs.insert(refSelf);
-    else
-      unboundRefs.insert(refSelf);
-  }
-
-  for (DIRecursiveTypeAttrInterface ref : boundRefs)
-    entry.pendingReplacements.erase(ref);
+  DependentTranslation &entry = cacheIter->second;
+  if (llvm::set_is_subset(entry.unboundSelfRefs,
+                          translationStack.back().second.unboundSelfRefs))
+    return entry.attr;
 
-  return std::make_pair(result, unboundRefs);
+  // Stale cache entry.
+  dependentCache.erase(cacheIter);
+  return {};
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/DebugImporter.h

@@ -116,12 +116,11 @@ class DebugImporter {
   public:
     RecursionPruner(MLIRContext *context) : context(context) {}
 
-    /// If this node was previously cached, returns the cached result.
     /// If this node is a recursive instance that was previously seen, returns a
-    /// self-reference.
-    /// Otherwise, returns null attr, and a translation stack frame is created
-    /// for this node. Expects `finalizeTranslation` & `popTranslationStack`
-    /// to be called on this node later.
+    /// self-reference. If this node was previously cached, returns the cached
+    /// result. Otherwise, returns null attr, and a translation stack frame is
+    /// created for this node. Expects `finalizeTranslation` &
+    /// `popTranslationStack` to be called on this node later.
     DINodeAttr pruneOrPushTranslationStack(llvm::DINode *node);
 
     /// Register the translated result of `node`. Returns the finalized result
@@ -135,33 +134,27 @@ class DebugImporter {
     void popTranslationStack(llvm::DINode *node);
 
   private:
-    /// Returns the cached result (if exists) with all known replacements
-    /// applied. Also returns the set of unbound self-refs that are unresolved
-    /// in this cached result. The cache entry will also be updated with the
-    /// replaced result and with the applied replacements removed from the
-    /// pendingReplacements map.
-    std::pair<DINodeAttr, DenseSet<DIRecursiveTypeAttrInterface>>
-    lookup(llvm::DINode *node);
+    /// Returns the cached result (if exists) or null.
+    /// The cache entry will be removed if not all of its dependent self-refs
+    /// exists.
+    DINodeAttr lookup(llvm::DINode *node);
 
     MLIRContext *context;
 
-    /// A lazy cached translation that contains the translated attribute as well
-    /// as any unbound self-references that need to be replaced at lookup.
-    struct CachedTranslation {
+    /// A cached translation that contains the translated attribute as well
+    /// as any unbound self-references that it depends on.
+    struct DependentTranslation {
       /// The translated attr. May contain unbound self-references for other
       /// recursive attrs.
       DINodeAttr attr;
-      /// The pending replacements that need to be run on this `attr` before it
-      /// can be used.
-      /// Each key is a recursive self-ref, and each value is a recursive decl
-      /// that may contain additional unbound self-refs that need to be
-      /// replaced. Each replacement will be applied at most once. Once a
-      /// replacement is applied, the cached `attr` will be updated, and the
-      /// replacement will be removed from this map.
-      DenseMap<DIRecursiveTypeAttrInterface, DINodeAttr> pendingReplacements;
+      /// The set of unbound self-refs that this cached entry refers to. All
+      /// these self-refs must exist for the cached entry to be valid.
+      DenseSet<DIRecursiveTypeAttrInterface> unboundSelfRefs;
     };
     /// A mapping between LLVM debug metadata and the corresponding attribute.
-    llvm::MapVector<llvm::DINode *, CachedTranslation> cache;
+    /// Only contains those with unboundSelfRefs. Fully self-contained attrs
+    /// will be cached by the outer main translator.
+    DenseMap<llvm::DINode *, DependentTranslation> dependentCache;
 
     /// Each potentially recursive node will have a TranslationState pushed onto
     /// the `translationStack` to keep track of whether this node is actually
@@ -171,9 +164,6 @@ class DebugImporter {
       /// instance of itself is seen while translating it). Null if this node
       /// has not been seen again deeper in the translation stack.
       DIRecursiveTypeAttrInterface recSelf;
-      /// The number of cache entries belonging to this layer of the translation
-      /// stack. This corresponds to the last `cacheSize` entries of `cache`.
-      uint64_t cacheSize = 0;
       /// All the unbound recursive self references in this layer of the
       /// translation stack.
       DenseSet<DIRecursiveTypeAttrInterface> unboundSelfRefs;
@@ -184,6 +174,11 @@ class DebugImporter {
     /// later when the stack is deeper, the node is recursive, and its
     /// TranslationState is assigned a recSelf.
     llvm::MapVector<llvm::DINode *, TranslationState> translationStack;
+
+    /// A mapping between DINodes that are recursive, and their assigned recId.
+    /// This is kept so that repeated occurrences of the same node can reuse the
+    /// same ID and be deduplicated.
+    DenseMap<llvm::DINode *, DistinctAttr> nodeToRecId;
   };
   RecursionPruner recursionPruner;