[MLIR][LLVM] DI Recursive Type fix for recursion via scope of composites

[zyx-billy]

Fixes this bug for the previous recursive DI type PR: #80251 (comment).

Drawing inspiration from how clang uses DIBuilder to build forward decls, this PR changes how placeholders are created & updated. Instead of requiring each recursive DIType to do in-place mutation, we simply ask for a temporary node as the placeholder, and run RAUW at the end when the concrete node is translated.

This has the side effect of simplifying what's needed to add recursion support for a type. Now only one additional method needs to be created for exporting. Concretely, for this PR, translateImpl for DICompositeType is back to the state it was before the previous PR, and the only net addition for DICompositeType is translateTemporaryImpl.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Billy Zhu (zyx-billy)

Changes

Fixes this bug for the previous recursive DI type PR: #80251 (comment).

Drawing inspiration from how clang uses DIBuilder to build forward decls, this PR changes how placeholders are created & updated. Instead of requiring each recursive DIType to do in-place mutation, we simply ask for a temporary node as the placeholder, and run RAUW at the end when the concrete node is translated.

This has the side effect of simplifying what's need to add recursion support for a type. Now only one additional method needs to be created for exporting. Concretely, for this PR, translateImpl for DICompositeType is back to the state it was before the previous PR, and the only net addition for DICompositeType is translateTemporaryImpl.

---

Full diff: https://github.com/llvm/llvm-project/pull/85850.diff

3 Files Affected:

• (modified) mlir/lib/Target/LLVMIR/DebugTranslation.cpp (+26-24)
• (modified) mlir/lib/Target/LLVMIR/DebugTranslation.h (+8-16)
• (modified) mlir/test/Target/LLVMIR/llvmir-debug.mlir (+21)

diff --git a/mlir/lib/Target/LLVMIR/DebugTranslation.cpp b/mlir/lib/Target/LLVMIR/DebugTranslation.cpp
index 1248de3db07048..126a671a58e808 100644
--- a/mlir/lib/Target/LLVMIR/DebugTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/DebugTranslation.cpp
@@ -116,8 +116,20 @@ static DINodeT *getDistinctOrUnique(bool isDistinct, Ts &&...args) {
   return DINodeT::get(std::forward<Ts>(args)...);
 }
 
+llvm::TempDICompositeType
+DebugTranslation::translateTemporaryImpl(DICompositeTypeAttr attr) {
+  return llvm::DICompositeType::getTemporary(
+      llvmCtx, attr.getTag(), getMDStringOrNull(attr.getName()), nullptr,
+      attr.getLine(), nullptr, nullptr, attr.getSizeInBits(),
+      attr.getAlignInBits(),
+      /*OffsetInBits=*/0,
+      /*Flags=*/static_cast<llvm::DINode::DIFlags>(attr.getFlags()),
+      /*Elements=*/nullptr, /*RuntimeLang=*/0,
+      /*VTableHolder=*/nullptr);
+}
+
 llvm::DICompositeType *
-DebugTranslation::translateImplGetPlaceholder(DICompositeTypeAttr attr) {
+DebugTranslation::translateImpl(DICompositeTypeAttr attr) {
   // TODO: Use distinct attributes to model this, once they have landed.
   // Depending on the tag, composite types must be distinct.
   bool isDistinct = false;
@@ -129,8 +141,10 @@ DebugTranslation::translateImplGetPlaceholder(DICompositeTypeAttr attr) {
     isDistinct = true;
   }
 
-  llvm::TempMDTuple placeholderElements =
-      llvm::MDNode::getTemporary(llvmCtx, std::nullopt);
+  SmallVector<llvm::Metadata *> elements;
+  for (DINodeAttr member : attr.getElements())
+    elements.push_back(translate(member));
+
   return getDistinctOrUnique<llvm::DICompositeType>(
       isDistinct, llvmCtx, attr.getTag(), getMDStringOrNull(attr.getName()),
       translate(attr.getFile()), attr.getLine(), translate(attr.getScope()),
@@ -138,23 +152,8 @@ DebugTranslation::translateImplGetPlaceholder(DICompositeTypeAttr attr) {
       attr.getAlignInBits(),
       /*OffsetInBits=*/0,
       /*Flags=*/static_cast<llvm::DINode::DIFlags>(attr.getFlags()),
-      /*Elements=*/placeholderElements.get(), /*RuntimeLang=*/0,
-      /*VTableHolder=*/nullptr);
-}
-
-void DebugTranslation::translateImplFillPlaceholder(
-    DICompositeTypeAttr attr, llvm::DICompositeType *placeholder) {
-  SmallVector<llvm::Metadata *> elements;
-  for (DINodeAttr member : attr.getElements())
-    elements.push_back(translate(member));
-  placeholder->replaceElements(llvm::MDNode::get(llvmCtx, elements));
-}
-
-llvm::DICompositeType *
-DebugTranslation::translateImpl(DICompositeTypeAttr attr) {
-  llvm::DICompositeType *placeholder = translateImplGetPlaceholder(attr);
-  translateImplFillPlaceholder(attr, placeholder);
-  return placeholder;
+      llvm::MDNode::get(llvmCtx, elements),
+      /*RuntimeLang=*/0, /*VTableHolder=*/nullptr);
 }
 
 llvm::DIDerivedType *DebugTranslation::translateImpl(DIDerivedTypeAttr attr) {
@@ -234,10 +233,13 @@ DebugTranslation::translateRecursive(DIRecursiveTypeAttrInterface attr) {
   llvm::DIType *result =
       TypeSwitch<DIRecursiveTypeAttrInterface, llvm::DIType *>(attr)
           .Case<DICompositeTypeAttr>([&](auto attr) {
-            auto *placeholder = translateImplGetPlaceholder(attr);
-            setRecursivePlaceholder(placeholder);
-            translateImplFillPlaceholder(attr, placeholder);
-            return placeholder;
+            auto temporary = translateTemporaryImpl(attr);
+            setRecursivePlaceholder(temporary.get());
+            // Must call `translateImpl` directly instead of `translate` to
+            // avoid handling the recursive interface again.
+            auto *concrete = translateImpl(attr);
+            temporary->replaceAllUsesWith(concrete);
+            return concrete;
           });
 
   assert(recursiveTypeMap.back().first == recursiveId &&
diff --git a/mlir/lib/Target/LLVMIR/DebugTranslation.h b/mlir/lib/Target/LLVMIR/DebugTranslation.h
index d2171fed8c594f..c7a5228cbf5e92 100644
--- a/mlir/lib/Target/LLVMIR/DebugTranslation.h
+++ b/mlir/lib/Target/LLVMIR/DebugTranslation.h
@@ -88,23 +88,16 @@ class DebugTranslation {
   llvm::DISubroutineType *translateImpl(DISubroutineTypeAttr attr);
   llvm::DIType *translateImpl(DITypeAttr attr);
 
-  /// Attributes that support self recursion need to implement two methods and
-  /// hook into the `translateImpl` overload for `DIRecursiveTypeAttr`.
-  /// - `<llvm type> translateImplGetPlaceholder(<mlir type>)`:
-  ///   Translate the DI attr without translating any potentially recursive
-  ///   nested DI attrs.
-  /// - `void translateImplFillPlaceholder(<mlir type>, <llvm type>)`:
-  ///   Given the placeholder returned by `translateImplGetPlaceholder`, fill
-  ///   any holes by recursively translating nested DI attrs. This method must
-  ///   mutate the placeholder that is passed in, instead of creating a new one.
+  /// Attributes that support self recursion need to implement an additional
+  /// method to hook into `translateRecursive`.
+  /// - `<temp llvm type> translateTemporaryImpl(<mlir type>)`:
+  ///   Create a temporary translation of the DI attr without recursively
+  ///   translating any nested DI attrs.
   llvm::DIType *translateRecursive(DIRecursiveTypeAttrInterface attr);
 
-  /// Get a placeholder DICompositeType without recursing into the elements.
-  llvm::DICompositeType *translateImplGetPlaceholder(DICompositeTypeAttr attr);
-  /// Completes the DICompositeType `placeholder` by recursively translating the
-  /// elements.
-  void translateImplFillPlaceholder(DICompositeTypeAttr attr,
-                                    llvm::DICompositeType *placeholder);
+  /// Translate the given attribute to a temporary llvm debug metadata of the
+  /// corresponding type.
+  llvm::TempDICompositeType translateTemporaryImpl(DICompositeTypeAttr attr);
 
   /// Constructs a string metadata node from the string attribute. Returns
   /// nullptr if `stringAttr` is null or contains and empty string.
@@ -121,7 +114,6 @@ class DebugTranslation {
   DenseMap<Attribute, llvm::DINode *> attrToNode;
 
   /// A mapping between recursive ID and the translated DIType.
-  /// DIType.
   llvm::MapVector<DistinctAttr, llvm::DIType *> recursiveTypeMap;
 
   /// A mapping between a distinct ID and the translated LLVM metadata node.
diff --git a/mlir/test/Target/LLVMIR/llvmir-debug.mlir b/mlir/test/Target/LLVMIR/llvmir-debug.mlir
index 5d70bff52bb2b9..33faa95d7d1967 100644
--- a/mlir/test/Target/LLVMIR/llvmir-debug.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-debug.mlir
@@ -402,3 +402,24 @@ llvm.func @func_debug_directives() {
 llvm.func @class_method() {
   llvm.return loc(#loc3)
 } loc(#loc3)
+
+// -----
+
+// Test recursive composite type with recursion via its scope is translated.
+
+#di_composite_type_self = #llvm.di_composite_type<tag = DW_TAG_null, recId = distinct[0]<>>
+#di_file = #llvm.di_file<"test.mlir" in "/">
+#di_subroutine_type = #llvm.di_subroutine_type<types = #di_composite_type_self>
+#di_subprogram = #llvm.di_subprogram<scope = #di_file, name = "my_func", file = #di_file, line = 182, scopeLine = 182, subprogramFlags = Optimized, type = #di_subroutine_type>
+// The composite type whose scope leads to a self-recursive cycle.
+#di_composite_type = #llvm.di_composite_type<tag = DW_TAG_class_type, recId = distinct[0]<>, name = "id", file = #di_file, line = 204, scope = #di_subprogram, flags = "Public|TypePassByReference|NonTrivial", sizeInBits = 128>
+#di_global_variable = #llvm.di_global_variable<scope = #di_file, name = "my_global", file = #di_file, line = 4641, type = #di_composite_type, isDefined = true>
+#di_global_variable_expression = #llvm.di_global_variable_expression<var = #di_global_variable, expr = <>>
+
+llvm.mlir.global @global_variable() {dbg_expr = #di_global_variable_expression} : !llvm.struct<(struct<(i64)>, i32)>
+
+// CHECK: distinct !DIGlobalVariable(name: "my_global", {{.*}}, type: ![[COMP:[0-9]+]],
+// CHECK: ![[COMP]] = distinct !DICompositeType(tag: DW_TAG_class_type, name: "id", scope: ![[SCOPE:[0-9]+]],
+// CHECK: ![[SCOPE]] = !DISubprogram(name: "my_func", {{.*}}, type: ![[SUBROUTINE:[0-9]+]],
+// CHECK: ![[SUBROUTINE]] = !DISubroutineType(types: ![[SR_TYPES:[0-9]+]])
+// CHECK: ![[SR_TYPES]] = !{![[COMP]]}

[gysit]

Thanks, LGTM!

I did not know that particular API but that is a nice simplification. If you see some opportunities to strip and/or format the test case a bit that would be great.

[zyx-billy]

Oh yes I can probably strip some fields from these attrs in the test case. Thanks!

[zyx-billy]

This windows build is taking forever... I'll check back tonight too but feel free to merge it if you need it sooner.