Skip to content

Commit dd86604

Browse files
zyx-billyzyx-billy
authored
[MLIR][LLVM] Use CyclicReplacerCache for recursive DIType import (#98203)
Use the new CyclicReplacerCache from #98202 to support importing of recursive DITypes in LLVM dialect's DebugImporter. This helps simplify the implementation, allows for separate testing of the cache infra itself, and as a result we even got more efficient translations.
1 parent 06bbbf1 commit dd86604

File tree

3 files changed

+51
-227
lines changed

3 files changed

+51
-227
lines changed

‎mlir/lib/Target/LLVMIR/DebugImporter.cpp

Lines changed: 26 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ using namespace mlir::LLVM::detail;
2828

2929
DebugImporter::DebugImporter(ModuleOp mlirModule,
3030
bool dropDICompositeTypeElements)
31-
: recursionPruner(mlirModule.getContext()),
31+
: cache([&](llvm::DINode *node) { return createRecSelf(node); }),
3232
context(mlirModule.getContext()), mlirModule(mlirModule),
3333
dropDICompositeTypeElements(dropDICompositeTypeElements) {}
3434

@@ -287,16 +287,9 @@ DINodeAttr DebugImporter::translate(llvm::DINode *node) {
287287
return nullptr;
288288

289289
// Check for a cached instance.
290-
if (DINodeAttr attr = nodeToAttr.lookup(node))
291-
return attr;
292-
293-
// Register with the recursive translator. If it can be handled without
294-
// recursing into it, return the result immediately.
295-
if (DINodeAttr attr = recursionPruner.pruneOrPushTranslationStack(node))
296-
return attr;
297-
298-
auto guard = llvm::make_scope_exit(
299-
[&]() { recursionPruner.popTranslationStack(node); });
290+
auto cacheEntry = cache.lookupOrInit(node);
291+
if (std::optional<DINodeAttr> result = cacheEntry.get())
292+
return *result;
300293

301294
// Convert the debug metadata if possible.
302295
auto translateNode = [this](llvm::DINode *node) -> DINodeAttr {
@@ -335,20 +328,20 @@ DINodeAttr DebugImporter::translate(llvm::DINode *node) {
335328
return nullptr;
336329
};
337330
if (DINodeAttr attr = translateNode(node)) {
338-
auto [result, isSelfContained] =
339-
recursionPruner.finalizeTranslation(node, attr);
340-
// Only cache fully self-contained nodes.
341-
if (isSelfContained)
342-
nodeToAttr.try_emplace(node, result);
343-
return result;
331+
// If this node was repeated, lookup its recursive ID and assign it to the
332+
// base result.
333+
if (cacheEntry.wasRepeated()) {
334+
DistinctAttr recId = nodeToRecId.lookup(node);
335+
auto recType = cast<DIRecursiveTypeAttrInterface>(attr);
336+
attr = cast<DINodeAttr>(recType.withRecId(recId));
337+
}
338+
cacheEntry.resolve(attr);
339+
return attr;
344340
}
341+
cacheEntry.resolve(nullptr);
345342
return nullptr;
346343
}
347344

348-
//===----------------------------------------------------------------------===//
349-
// RecursionPruner
350-
//===----------------------------------------------------------------------===//
351-
352345
/// Get the `getRecSelf` constructor for the translated type of `node` if its
353346
/// translated DITypeAttr supports recursion. Otherwise, returns nullptr.
354347
static function_ref<DIRecursiveTypeAttrInterface(DistinctAttr)>
@@ -361,104 +354,20 @@ getRecSelfConstructor(llvm::DINode *node) {
361354
.Default(CtorType());
362355
}
363356

364-
DINodeAttr DebugImporter::RecursionPruner::pruneOrPushTranslationStack(
365-
llvm::DINode *node) {
366-
// If the node type is capable of being recursive, check if it's seen
367-
// before.
357+
std::optional<DINodeAttr> DebugImporter::createRecSelf(llvm::DINode *node) {
368358
auto recSelfCtor = getRecSelfConstructor(node);
369-
if (recSelfCtor) {
370-
// If a cyclic dependency is detected since the same node is being
371-
// traversed twice, emit a recursive self type, and mark the duplicate
372-
// node on the translationStack so it can emit a recursive decl type.
373-
auto [iter, inserted] = translationStack.try_emplace(node);
374-
if (!inserted) {
375-
// The original node may have already been assigned a recursive ID from
376-
// a different self-reference. Use that if possible.
377-
DIRecursiveTypeAttrInterface recSelf = iter->second.recSelf;
378-
if (!recSelf) {
379-
DistinctAttr recId = nodeToRecId.lookup(node);
380-
if (!recId) {
381-
recId = DistinctAttr::create(UnitAttr::get(context));
382-
nodeToRecId[node] = recId;
383-
}
384-
recSelf = recSelfCtor(recId);
385-
iter->second.recSelf = recSelf;
386-
}
387-
// Inject the self-ref into the previous layer.
388-
translationStack.back().second.unboundSelfRefs.insert(recSelf);
389-
return cast<DINodeAttr>(recSelf);
390-
}
359+
if (!recSelfCtor)
360+
return std::nullopt;
361+
362+
// The original node may have already been assigned a recursive ID from
363+
// a different self-reference. Use that if possible.
364+
DistinctAttr recId = nodeToRecId.lookup(node);
365+
if (!recId) {
366+
recId = DistinctAttr::create(UnitAttr::get(context));
367+
nodeToRecId[node] = recId;
391368
}
392-
393-
return lookup(node);
394-
}
395-
396-
std::pair<DINodeAttr, bool>
397-
DebugImporter::RecursionPruner::finalizeTranslation(llvm::DINode *node,
398-
DINodeAttr result) {
399-
// If `node` is not a potentially recursive type, it will not be on the
400-
// translation stack. Nothing to set in this case.
401-
if (translationStack.empty())
402-
return {result, true};
403-
if (translationStack.back().first != node)
404-
return {result, translationStack.back().second.unboundSelfRefs.empty()};
405-
406-
TranslationState &state = translationStack.back().second;
407-
408-
// If this node is actually recursive, set the recId onto `result`.
409-
if (DIRecursiveTypeAttrInterface recSelf = state.recSelf) {
410-
auto recType = cast<DIRecursiveTypeAttrInterface>(result);
411-
result = cast<DINodeAttr>(recType.withRecId(recSelf.getRecId()));
412-
// Remove this recSelf from the set of unbound selfRefs.
413-
state.unboundSelfRefs.erase(recSelf);
414-
}
415-
416-
// Insert the result into our internal cache if it's not self-contained.
417-
if (!state.unboundSelfRefs.empty()) {
418-
[[maybe_unused]] auto [_, inserted] = dependentCache.try_emplace(
419-
node, DependentTranslation{result, state.unboundSelfRefs});
420-
assert(inserted && "invalid state: caching the same DINode twice");
421-
return {result, false};
422-
}
423-
return {result, true};
424-
}
425-
426-
void DebugImporter::RecursionPruner::popTranslationStack(llvm::DINode *node) {
427-
// If `node` is not a potentially recursive type, it will not be on the
428-
// translation stack. Nothing to handle in this case.
429-
if (translationStack.empty() || translationStack.back().first != node)
430-
return;
431-
432-
// At the end of the stack, all unbound self-refs must be resolved already,
433-
// and the entire cache should be accounted for.
434-
TranslationState &currLayerState = translationStack.back().second;
435-
if (translationStack.size() == 1) {
436-
assert(currLayerState.unboundSelfRefs.empty() &&
437-
"internal error: unbound recursive self reference at top level.");
438-
translationStack.pop_back();
439-
return;
440-
}
441-
442-
// Copy unboundSelfRefs down to the previous level.
443-
TranslationState &nextLayerState = (++translationStack.rbegin())->second;
444-
nextLayerState.unboundSelfRefs.insert(currLayerState.unboundSelfRefs.begin(),
445-
currLayerState.unboundSelfRefs.end());
446-
translationStack.pop_back();
447-
}
448-
449-
DINodeAttr DebugImporter::RecursionPruner::lookup(llvm::DINode *node) {
450-
auto cacheIter = dependentCache.find(node);
451-
if (cacheIter == dependentCache.end())
452-
return {};
453-
454-
DependentTranslation &entry = cacheIter->second;
455-
if (llvm::set_is_subset(entry.unboundSelfRefs,
456-
translationStack.back().second.unboundSelfRefs))
457-
return entry.attr;
458-
459-
// Stale cache entry.
460-
dependentCache.erase(cacheIter);
461-
return {};
369+
DIRecursiveTypeAttrInterface recSelf = recSelfCtor(recId);
370+
return cast<DINodeAttr>(recSelf);
462371
}
463372

464373
//===----------------------------------------------------------------------===//

‎mlir/lib/Target/LLVMIR/DebugImporter.h

Lines changed: 9 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
1818
#include "mlir/IR/BuiltinOps.h"
1919
#include "mlir/IR/MLIRContext.h"
20+
#include "mlir/Support/CyclicReplacerCache.h"
2021
#include "llvm/IR/DebugInfoMetadata.h"
2122

2223
namespace mlir {
@@ -87,102 +88,18 @@ class DebugImporter {
8788
/// for it, or create a new one if not.
8889
DistinctAttr getOrCreateDistinctID(llvm::DINode *node);
8990

90-
/// A mapping between LLVM debug metadata and the corresponding attribute.
91-
DenseMap<llvm::DINode *, DINodeAttr> nodeToAttr;
91+
std::optional<DINodeAttr> createRecSelf(llvm::DINode *node);
92+
9293
/// A mapping between distinct LLVM debug metadata nodes and the corresponding
9394
/// distinct id attribute.
9495
DenseMap<llvm::DINode *, DistinctAttr> nodeToDistinctAttr;
9596

96-
/// Translation helper for recursive DINodes.
97-
/// Works alongside a stack-based DINode translator (the "main translator")
98-
/// for gracefully handling DINodes that are recursive.
99-
///
100-
/// Usage:
101-
/// - Before translating a node, call `pruneOrPushTranslationStack` to see if
102-
/// the pruner can preempt this translation. If this is a node that the
103-
/// pruner already knows how to handle, it will return the translated
104-
/// DINodeAttr.
105-
/// - After a node is successfully translated by the main translator, call
106-
/// `finalizeTranslation` to save the translated result with the pruner, and
107-
/// give it a chance to further modify the result.
108-
/// - Regardless of success or failure by the main translator, always call
109-
/// `popTranslationStack` at the end of translating a node. This is
110-
/// necessary to keep the internal book-keeping in sync.
111-
///
112-
/// This helper maintains an internal cache so that no recursive type will
113-
/// be translated more than once by the main translator.
114-
/// This internal cache is different from the cache maintained by the main
115-
/// translator because it may store nodes that are not self-contained (i.e.
116-
/// contain unbounded recursive self-references).
117-
class RecursionPruner {
118-
public:
119-
RecursionPruner(MLIRContext *context) : context(context) {}
120-
121-
/// If this node is a recursive instance that was previously seen, returns a
122-
/// self-reference. If this node was previously cached, returns the cached
123-
/// result. Otherwise, returns null attr, and a translation stack frame is
124-
/// created for this node. Expects `finalizeTranslation` &
125-
/// `popTranslationStack` to be called on this node later.
126-
DINodeAttr pruneOrPushTranslationStack(llvm::DINode *node);
127-
128-
/// Register the translated result of `node`. Returns the finalized result
129-
/// (with recId if recursive) and whether the result is self-contained
130-
/// (i.e. contains no unbound self-refs).
131-
std::pair<DINodeAttr, bool> finalizeTranslation(llvm::DINode *node,
132-
DINodeAttr result);
133-
134-
/// Pop off a frame from the translation stack after a node is done being
135-
/// translated.
136-
void popTranslationStack(llvm::DINode *node);
137-
138-
private:
139-
/// Returns the cached result (if exists) or null.
140-
/// The cache entry will be removed if not all of its dependent self-refs
141-
/// exists.
142-
DINodeAttr lookup(llvm::DINode *node);
143-
144-
MLIRContext *context;
145-
146-
/// A cached translation that contains the translated attribute as well
147-
/// as any unbound self-references that it depends on.
148-
struct DependentTranslation {
149-
/// The translated attr. May contain unbound self-references for other
150-
/// recursive attrs.
151-
DINodeAttr attr;
152-
/// The set of unbound self-refs that this cached entry refers to. All
153-
/// these self-refs must exist for the cached entry to be valid.
154-
DenseSet<DIRecursiveTypeAttrInterface> unboundSelfRefs;
155-
};
156-
/// A mapping between LLVM debug metadata and the corresponding attribute.
157-
/// Only contains those with unboundSelfRefs. Fully self-contained attrs
158-
/// will be cached by the outer main translator.
159-
DenseMap<llvm::DINode *, DependentTranslation> dependentCache;
160-
161-
/// Each potentially recursive node will have a TranslationState pushed onto
162-
/// the `translationStack` to keep track of whether this node is actually
163-
/// recursive (i.e. has self-references inside), and other book-keeping.
164-
struct TranslationState {
165-
/// The rec-self if this node is indeed a recursive node (i.e. another
166-
/// instance of itself is seen while translating it). Null if this node
167-
/// has not been seen again deeper in the translation stack.
168-
DIRecursiveTypeAttrInterface recSelf;
169-
/// All the unbound recursive self references in this layer of the
170-
/// translation stack.
171-
DenseSet<DIRecursiveTypeAttrInterface> unboundSelfRefs;
172-
};
173-
/// A stack that stores the metadata nodes that are being traversed. The
174-
/// stack is used to handle cyclic dependencies during metadata translation.
175-
/// Each node is pushed with an empty TranslationState. If it is ever seen
176-
/// later when the stack is deeper, the node is recursive, and its
177-
/// TranslationState is assigned a recSelf.
178-
llvm::MapVector<llvm::DINode *, TranslationState> translationStack;
179-
180-
/// A mapping between DINodes that are recursive, and their assigned recId.
181-
/// This is kept so that repeated occurrences of the same node can reuse the
182-
/// same ID and be deduplicated.
183-
DenseMap<llvm::DINode *, DistinctAttr> nodeToRecId;
184-
};
185-
RecursionPruner recursionPruner;
97+
/// A mapping between DINodes that are recursive, and their assigned recId.
98+
/// This is kept so that repeated occurrences of the same node can reuse the
99+
/// same ID and be deduplicated.
100+
DenseMap<llvm::DINode *, DistinctAttr> nodeToRecId;
101+
102+
CyclicReplacerCache<llvm::DINode *, DINodeAttr> cache;
186103

187104
MLIRContext *context;
188105
ModuleOp mlirModule;

‎mlir/test/Target/LLVMIR/Import/debug-info.ll

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -308,16 +308,16 @@ define void @class_method() {
308308
}
309309

310310
; Verify the cyclic composite type is identified, even though conversion begins from the subprogram type.
311-
; CHECK: #[[COMP_SELF:.+]] = #llvm.di_composite_type<tag = DW_TAG_null, recId = [[REC_ID:.+]]>
312-
; CHECK: #[[COMP_PTR:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP_SELF]], sizeInBits = 64>
313-
; CHECK: #[[SP_TYPE:.+]] = #llvm.di_subroutine_type<types = #{{.*}}, #[[COMP_PTR]]>
314-
; CHECK: #[[SP_INNER:.+]] = #llvm.di_subprogram<id = [[SP_ID:.+]], compileUnit = #{{.*}}, scope = #[[COMP_SELF]], name = "class_method", file = #{{.*}}, subprogramFlags = Definition, type = #[[SP_TYPE]]>
315-
; CHECK: #[[COMP:.+]] = #llvm.di_composite_type<tag = DW_TAG_class_type, recId = [[REC_ID]], name = "class_name", file = #{{.*}}, line = 42, flags = "TypePassByReference|NonTrivial", elements = #[[SP_INNER]]>
311+
; CHECK-DAG: #[[COMP_SELF:.+]] = #llvm.di_composite_type<tag = DW_TAG_null, recId = [[REC_ID:.+]]>
312+
; CHECK-DAG: #[[COMP_PTR:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP_SELF]], sizeInBits = 64>
313+
; CHECK-DAG: #[[SP_TYPE:.+]] = #llvm.di_subroutine_type<types = #{{.*}}, #[[COMP_PTR]]>
314+
; CHECK-DAG: #[[SP_INNER:.+]] = #llvm.di_subprogram<id = [[SP_ID:.+]], compileUnit = #{{.*}}, scope = #[[COMP_SELF]], name = "class_method", file = #{{.*}}, subprogramFlags = Definition, type = #[[SP_TYPE]]>
315+
; CHECK-DAG: #[[COMP:.+]] = #llvm.di_composite_type<tag = DW_TAG_class_type, recId = [[REC_ID]], name = "class_name", file = #{{.*}}, line = 42, flags = "TypePassByReference|NonTrivial", elements = #[[SP_INNER]]>
316316

317-
; CHECK: #[[COMP_PTR_OUTER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP]], sizeInBits = 64>
318-
; CHECK: #[[SP_TYPE_OUTER:.+]] = #llvm.di_subroutine_type<types = #{{.*}}, #[[COMP_PTR_OUTER]]>
319-
; CHECK: #[[SP_OUTER:.+]] = #llvm.di_subprogram<id = [[SP_ID]], compileUnit = #{{.*}}, scope = #[[COMP]], name = "class_method", file = #{{.*}}, subprogramFlags = Definition, type = #[[SP_TYPE_OUTER]]>
320-
; CHECK: #[[LOC]] = loc(fused<#[[SP_OUTER]]>
317+
; CHECK-DAG: #[[COMP_PTR_OUTER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP]], sizeInBits = 64>
318+
; CHECK-DAG: #[[SP_TYPE_OUTER:.+]] = #llvm.di_subroutine_type<types = #{{.*}}, #[[COMP_PTR_OUTER]]>
319+
; CHECK-DAG: #[[SP_OUTER:.+]] = #llvm.di_subprogram<id = [[SP_ID]], compileUnit = #{{.*}}, scope = #[[COMP]], name = "class_method", file = #{{.*}}, subprogramFlags = Definition, type = #[[SP_TYPE_OUTER]]>
320+
; CHECK-DAG: #[[LOC]] = loc(fused<#[[SP_OUTER]]>
321321

322322
!llvm.dbg.cu = !{!1}
323323
!llvm.module.flags = !{!0}
@@ -335,12 +335,12 @@ define void @class_method() {
335335
; // -----
336336

337337
; Verify the cyclic composite type is handled correctly.
338-
; CHECK: #[[COMP_SELF:.+]] = #llvm.di_composite_type<tag = DW_TAG_null, recId = [[REC_ID:.+]]>
339-
; CHECK: #[[COMP_PTR_INNER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP_SELF]]>
340-
; CHECK: #[[FIELD:.+]] = #llvm.di_derived_type<tag = DW_TAG_member, name = "call_field", baseType = #[[COMP_PTR_INNER]]>
341-
; CHECK: #[[COMP:.+]] = #llvm.di_composite_type<tag = DW_TAG_class_type, recId = [[REC_ID]], name = "class_field", file = #{{.*}}, line = 42, flags = "TypePassByReference|NonTrivial", elements = #[[FIELD]]>
342-
; CHECK: #[[COMP_PTR_OUTER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP]]>
343-
; CHECK: #[[VAR0:.+]] = #llvm.di_local_variable<scope = #{{.*}}, name = "class_field", file = #{{.*}}, type = #[[COMP_PTR_OUTER]]>
338+
; CHECK-DAG: #[[COMP_SELF:.+]] = #llvm.di_composite_type<tag = DW_TAG_null, recId = [[REC_ID:.+]]>
339+
; CHECK-DAG: #[[COMP_PTR_INNER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP_SELF]]>
340+
; CHECK-DAG: #[[FIELD:.+]] = #llvm.di_derived_type<tag = DW_TAG_member, name = "call_field", baseType = #[[COMP_PTR_INNER]]>
341+
; CHECK-DAG: #[[COMP:.+]] = #llvm.di_composite_type<tag = DW_TAG_class_type, recId = [[REC_ID]], name = "class_field", file = #{{.*}}, line = 42, flags = "TypePassByReference|NonTrivial", elements = #[[FIELD]]>
342+
; CHECK-DAG: #[[COMP_PTR_OUTER:.+]] = #llvm.di_derived_type<tag = DW_TAG_pointer_type, baseType = #[[COMP]]>
343+
; CHECK-DAG: #[[VAR0:.+]] = #llvm.di_local_variable<scope = #{{.*}}, name = "class_field", file = #{{.*}}, type = #[[COMP_PTR_OUTER]]>
344344

345345
; CHECK: @class_field
346346
; CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
@@ -727,9 +727,7 @@ define void @class_field(ptr %arg1) !dbg !18 {
727727
; CHECK-DAG: #[[B_TO_C]] = #llvm.di_derived_type<{{.*}}name = "->C", {{.*}}baseType = #[[C_FROM_B:.+]]>
728728
; CHECK-DAG: #[[C_FROM_B]] = #llvm.di_composite_type<{{.*}}recId = [[C_RECID:.+]], {{.*}}name = "C", {{.*}}elements = #[[TO_A_SELF:.+]], #[[TO_B_SELF:.+]], #[[TO_C_SELF:.+]]>
729729

730-
; CHECK-DAG: #[[C_FROM_A]] = #llvm.di_composite_type<{{.*}}recId = [[C_RECID]], {{.*}}name = "C", {{.*}}elements = #[[TO_A_SELF]], #[[TO_B_INNER:.+]], #[[TO_C_SELF]]
731-
; CHECK-DAG: #[[TO_B_INNER]] = #llvm.di_derived_type<{{.*}}name = "->B", {{.*}}baseType = #[[B_INNER:.+]]>
732-
; CHECK-DAG: #[[B_INNER]] = #llvm.di_composite_type<{{.*}}name = "B", {{.*}}elements = #[[TO_C_SELF]]>
730+
; CHECK-DAG: #[[C_FROM_A]] = #llvm.di_composite_type<{{.*}}recId = [[C_RECID]], {{.*}}name = "C", {{.*}}elements = #[[TO_A_SELF]], #[[A_TO_B:.+]], #[[TO_C_SELF]]
733731

734732
; CHECK-DAG: #[[TO_A_SELF]] = #llvm.di_derived_type<{{.*}}name = "->A", {{.*}}baseType = #[[A_SELF:.+]]>
735733
; CHECK-DAG: #[[TO_B_SELF]] = #llvm.di_derived_type<{{.*}}name = "->B", {{.*}}baseType = #[[B_SELF:.+]]>

0 commit comments

Comments
 (0)