generic recursive type with importer impl

zyx-billy · zyx-billy · commit 4c03e15c1222 · 2024-02-27T14:32:03.000-08:00
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -301,6 +301,70 @@ def LLVM_DIExpressionAttr : LLVM_Attr<"DIExpression", "di_expression"> {
   let assemblyFormat = "`<` ( `[` $operations^ `]` ) : (``)? `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// DIRecursiveTypeAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_DIRecursiveTypeAttr : LLVM_Attr<"DIRecursiveType", "di_recursive_type",
+                                    /*traits=*/[], "DITypeAttr"> {
+  let description = [{
+    This attribute enables recursive DITypes. There are two modes for this
+    attribute.
+
+    1. If `baseType` is present:
+      - This type is considered a recursive declaration (rec-decl).
+      - The `baseType` is a self-recursive type identified by the `id` field.
+
+    2. If `baseType` is not present:
+      - This type is considered a recursive self reference (rec-self).
+      - This DIRecursiveType itself is a placeholder type that should be
+        conceptually replaced with the closet parent DIRecursiveType with the
+        same `id` field.
+
+    e.g. To represent a linked list struct:
+
+      #rec_self = di_recursive_type<self_id = 0>
+      #ptr = di_derived_type<baseType: #rec_self, ...>
+      #field = di_derived_type<name = "next", baseType: #ptr, ...>
+      #struct = di_composite_type<name = "Node", elements: #field, ...>
+      #rec = di_recursive_type<self_id = 0, baseType: #struct>
+
+      #var = di_local_variable<type = #struct_type, ...>
+
+    Note that the a rec-self without an outer rec-decl with the same id is
+    conceptually the same as an "unbound" variable. The context needs to provide
+    meaning to the rec-self.
+
+    This can be avoided by calling the `getUnfoldedBaseType()` method on a
+    rec-decl, which returns the `baseType` with all matching rec-self instances
+    replaced with this rec-decl again. This is useful, for example, for fetching
+    a field out of a recursive struct and maintaining the legality of the field
+    type.
+  }];
+
+  let parameters = (ins
+    "DistinctAttr":$id,
+    OptionalParameter<"DITypeAttr">:$baseType
+  );
+
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "DistinctAttr":$id), [{
+      return $_get(id.getContext(), id, nullptr);
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Whether this node represents a self-reference.
+    bool isRecSelf() { return !getBaseType(); }
+
+    /// Get the `baseType` with all instances of the corresponding rec-self
+    /// replaced with this attribute. This can only be called if `!isRecSelf()`.
+    DITypeAttr getUnfoldedBaseType();
+  }];
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // DINullTypeAttr
 //===----------------------------------------------------------------------===//
@@ -526,14 +590,15 @@ def LLVM_DISubprogramAttr : LLVM_Attr<"DISubprogram", "di_subprogram",
     OptionalParameter<"unsigned">:$line,
     OptionalParameter<"unsigned">:$scopeLine,
     "DISubprogramFlags":$subprogramFlags,
-    OptionalParameter<"DISubroutineTypeAttr">:$type
+    OptionalParameter<"DIRecursiveTypeAttrOf<DISubroutineTypeAttr>">:$type
   );
   let builders = [
     AttrBuilderWithInferredContext<(ins
       "DistinctAttr":$id, "DICompileUnitAttr":$compileUnit,
       "DIScopeAttr":$scope, "StringRef":$name, "StringRef":$linkageName,
       "DIFileAttr":$file, "unsigned":$line, "unsigned":$scopeLine,
-      "DISubprogramFlags":$subprogramFlags, "DISubroutineTypeAttr":$type
+      "DISubprogramFlags":$subprogramFlags,
+      "DIRecursiveTypeAttrOf<DISubroutineTypeAttr>":$type
     ), [{
       MLIRContext *ctx = file.getContext();
       return $_get(ctx, id, compileUnit, scope, StringAttr::get(ctx, name),
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -52,9 +52,9 @@ class DILocalScopeAttr : public DIScopeAttr {
 };
 
 /// This class represents a LLVM attribute that describes a debug info type.
-class DITypeAttr : public DINodeAttr {
+class DITypeAttr : public DIScopeAttr {
 public:
-  using DINodeAttr::DINodeAttr;
+  using DIScopeAttr::DIScopeAttr;
 
   /// Support LLVM type casting.
   static bool classof(Attribute attr);
@@ -74,6 +74,10 @@ class TBAANodeAttr : public Attribute {
   }
 };
 
+// Forward declare.
+template <typename BaseType>
+class DIRecursiveTypeAttrOf;
+
 // Inline the LLVM generated Linkage enum and utility.
 // This is only necessary to isolate the "enum generated code" from the
 // attribute definition itself.
@@ -87,4 +91,31 @@ using linkage::Linkage;
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/LLVMIR/LLVMOpsAttrDefs.h.inc"
 
+namespace mlir {
+namespace LLVM {
+/// This class represents either a concrete attr, or a DIRecursiveTypeAttr
+/// containing such a concrete attr.
+template <typename BaseType>
+class DIRecursiveTypeAttrOf : public DITypeAttr {
+public:
+  static_assert(std::is_base_of_v<DITypeAttr, BaseType>);
+  using DITypeAttr::DITypeAttr;
+  /// Support LLVM type casting.
+  static bool classof(Attribute attr) {
+    if (auto rec = llvm::dyn_cast<DIRecursiveTypeAttr>(attr))
+      return llvm::isa<BaseType>(rec.getBaseType());
+    return llvm::isa<BaseType>(attr);
+  }
+
+  DIRecursiveTypeAttrOf(BaseType baseType) : DITypeAttr(baseType) {}
+
+  BaseType getUnfoldedBaseType() {
+    if (auto rec = llvm::dyn_cast<DIRecursiveTypeAttr>(this))
+      return llvm::cast<BaseType>(rec.getUnfoldedBaseType());
+    return llvm::cast<BaseType>(this);
+  }
+};
+} // namespace LLVM
+} // namespace mlir
+
 #endif // MLIR_DIALECT_LLVMIR_LLVMATTRS_H_
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -68,8 +68,8 @@ bool DINodeAttr::classof(Attribute attr) {
 //===----------------------------------------------------------------------===//
 
 bool DIScopeAttr::classof(Attribute attr) {
-  return llvm::isa<DICompileUnitAttr, DICompositeTypeAttr, DIFileAttr,
-                   DILocalScopeAttr, DIModuleAttr, DINamespaceAttr>(attr);
+  return llvm::isa<DICompileUnitAttr, DIFileAttr, DILocalScopeAttr,
+                   DIModuleAttr, DINamespaceAttr, DITypeAttr>(attr);
 }
 
 //===----------------------------------------------------------------------===//
@@ -86,8 +86,9 @@ bool DILocalScopeAttr::classof(Attribute attr) {
 //===----------------------------------------------------------------------===//
 
 bool DITypeAttr::classof(Attribute attr) {
-  return llvm::isa<DINullTypeAttr, DIBasicTypeAttr, DICompositeTypeAttr,
-                   DIDerivedTypeAttr, DISubroutineTypeAttr>(attr);
+  return llvm::isa<DIRecursiveTypeAttr, DINullTypeAttr, DIBasicTypeAttr,
+                   DICompositeTypeAttr, DIDerivedTypeAttr,
+                   DISubroutineTypeAttr>(attr);
 }
 
 //===----------------------------------------------------------------------===//
@@ -185,6 +186,19 @@ void printExpressionArg(AsmPrinter &printer, uint64_t opcode,
   });
 }
 
+//===----------------------------------------------------------------------===//
+// DIRecursiveTypeAttr
+//===----------------------------------------------------------------------===//
+DITypeAttr DIRecursiveTypeAttr::getUnfoldedBaseType() {
+  assert(!isRecSelf() && "cannot get baseType from a rec-self type");
+  return llvm::cast<DITypeAttr>(getBaseType().replace(
+      [&](DIRecursiveTypeAttr rec) -> std::optional<DIRecursiveTypeAttr> {
+        if (rec.getId() == getId())
+          return *this;
+        return std::nullopt;
+      }));
+}
+
 //===----------------------------------------------------------------------===//
 // TargetFeaturesAttr
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/DebugImporter.cpp b/mlir/lib/Target/LLVMIR/DebugImporter.cpp
@@ -51,10 +51,9 @@ DICompileUnitAttr DebugImporter::translateImpl(llvm::DICompileUnit *node) {
   std::optional<DIEmissionKind> emissionKind =
       symbolizeDIEmissionKind(node->getEmissionKind());
   return DICompileUnitAttr::get(
-      context, DistinctAttr::create(UnitAttr::get(context)),
-      node->getSourceLanguage(), translate(node->getFile()),
-      getStringAttrOrNull(node->getRawProducer()), node->isOptimized(),
-      emissionKind.value());
+      context, getOrCreateDistinctID(node), node->getSourceLanguage(),
+      translate(node->getFile()), getStringAttrOrNull(node->getRawProducer()),
+      node->isOptimized(), emissionKind.value());
 }
 
 DICompositeTypeAttr DebugImporter::translateImpl(llvm::DICompositeType *node) {
@@ -64,11 +63,7 @@ DICompositeTypeAttr DebugImporter::translateImpl(llvm::DICompositeType *node) {
     assert(element && "expected a non-null element type");
     elements.push_back(translate(element));
   }
-  // Drop the elements parameter if a cyclic dependency is detected. We
-  // currently cannot model these cycles and thus drop the parameter if
-  // required. A cyclic dependency is detected if one of the element nodes
-  // translates to a nullptr since the node is already on the translation stack.
-  // TODO: Support debug metadata with cyclic dependencies.
+  // Drop the elements parameter if any of the elements are invalid.
   if (llvm::is_contained(elements, nullptr))
     elements.clear();
   DITypeAttr baseType = translate(node->getBaseType());
@@ -84,7 +79,7 @@ DICompositeTypeAttr DebugImporter::translateImpl(llvm::DICompositeType *node) {
 }
 
 DIDerivedTypeAttr DebugImporter::translateImpl(llvm::DIDerivedType *node) {
-  // Return nullptr if the base type is a cyclic dependency.
+  // Return nullptr if the base type invalid.
   DITypeAttr baseType = translate(node->getBaseType());
   if (node->getBaseType() && !baseType)
     return nullptr;
@@ -179,14 +174,14 @@ DISubprogramAttr DebugImporter::translateImpl(llvm::DISubprogram *node) {
   // Only definitions require a distinct identifier.
   mlir::DistinctAttr id;
   if (node->isDistinct())
-    id = DistinctAttr::create(UnitAttr::get(context));
+    id = getOrCreateDistinctID(node);
   std::optional<DISubprogramFlags> subprogramFlags =
       symbolizeDISubprogramFlags(node->getSubprogram()->getSPFlags());
-  // Return nullptr if the scope or type is a cyclic dependency.
-  DIScopeAttr scope = translate(node->getScope());
+  // Return nullptr if the scope or type is invalid.
+  DIScopeAttr scope = cast<DIScopeAttr>(translate(node->getScope()));
   if (node->getScope() && !scope)
     return nullptr;
-  DISubroutineTypeAttr type = translate(node->getType());
+  DIRecursiveTypeAttrOf<DISubroutineTypeAttr> type = translate(node->getType());
   if (node->getType() && !type)
     return nullptr;
   return DISubprogramAttr::get(context, id, translate(node->getUnit()), scope,
@@ -229,7 +224,7 @@ DebugImporter::translateImpl(llvm::DISubroutineType *node) {
     }
     types.push_back(translate(type));
   }
-  // Return nullptr if any of the types is a cyclic dependency.
+  // Return nullptr if any of the types is invalid.
   if (llvm::is_contained(types, nullptr))
     return nullptr;
   return DISubroutineTypeAttr::get(context, node->getCC(), types);
@@ -247,12 +242,47 @@ DINodeAttr DebugImporter::translate(llvm::DINode *node) {
   if (DINodeAttr attr = nodeToAttr.lookup(node))
     return attr;
 
-  // Return nullptr if a cyclic dependency is detected since the same node is
-  // being traversed twice. This check avoids infinite recursion if the debug
-  // metadata contains cycles.
-  if (!translationStack.insert(node))
-    return nullptr;
-  auto guard = llvm::make_scope_exit([&]() { translationStack.pop_back(); });
+  // If a cyclic dependency is detected since the same node is being traversed
+  // twice, emit a recursive self type, and mark the duplicate node on the
+  // translationStack so it can emit a recursive decl type.
+  auto *typeNode = dyn_cast<llvm::DIType>(node);
+  if (typeNode) {
+    auto [iter, inserted] = typeTranslationStack.try_emplace(typeNode, nullptr);
+    if (!inserted) {
+      // The original node may have already been assigned a recursive ID from
+      // a different self-reference. Use that if possible.
+      DistinctAttr recId = iter->second;
+      if (!recId) {
+        recId = DistinctAttr::create(UnitAttr::get(context));
+        iter->second = recId;
+      }
+      unboundRecursiveSelfRefs.back().insert(recId);
+      return DIRecursiveTypeAttr::get(recId);
+    }
+  } else {
+    bool inserted =
+        nonTypeTranslationStack.insert({node, typeTranslationStack.size()});
+    assert(inserted && "recursion is only supported via DITypes");
+  }
+
+  unboundRecursiveSelfRefs.emplace_back();
+
+  auto guard = llvm::make_scope_exit([&]() {
+    if (typeNode)
+      typeTranslationStack.pop_back();
+    else
+      nonTypeTranslationStack.pop_back();
+
+    // Copy unboundRecursiveSelfRefs down to the previous level.
+    if (unboundRecursiveSelfRefs.size() == 1)
+      assert(unboundRecursiveSelfRefs.back().empty() &&
+             "internal error: unbound recursive self reference at top level.");
+    else
+      unboundRecursiveSelfRefs[unboundRecursiveSelfRefs.size() - 2].insert(
+          unboundRecursiveSelfRefs.back().begin(),
+          unboundRecursiveSelfRefs.back().end());
+    unboundRecursiveSelfRefs.pop_back();
+  });
 
   // Convert the debug metadata if possible.
   auto translateNode = [this](llvm::DINode *node) -> DINodeAttr {
@@ -289,7 +319,21 @@ DINodeAttr DebugImporter::translate(llvm::DINode *node) {
     return nullptr;
   };
   if (DINodeAttr attr = translateNode(node)) {
-    nodeToAttr.insert({node, attr});
+    // If this node was marked as recursive, wrap with a recursive type.
+    if (typeNode) {
+      if (DistinctAttr id = typeTranslationStack.lookup(typeNode)) {
+        DITypeAttr typeAttr = cast<DITypeAttr>(attr);
+        attr = DIRecursiveTypeAttr::get(context, id, typeAttr);
+
+        // Remove the unbound recursive attr.
+        AttrTypeReplacer replacer;
+        unboundRecursiveSelfRefs.back().erase(id);
+      }
+    }
+
+    // Only cache fully self-contained nodes.
+    if (unboundRecursiveSelfRefs.back().empty())
+      nodeToAttr.try_emplace(node, attr);
     return attr;
   }
   return nullptr;
@@ -346,3 +390,10 @@ StringAttr DebugImporter::getStringAttrOrNull(llvm::MDString *stringNode) {
     return StringAttr();
   return StringAttr::get(context, stringNode->getString());
 }
+
+DistinctAttr DebugImporter::getOrCreateDistinctID(llvm::DINode *node) {
+  DistinctAttr &id = nodeToDistinctAttr[node];
+  if (!id)
+    id = DistinctAttr::create(UnitAttr::get(context));
+  return id;
+}
diff --git a/mlir/lib/Target/LLVMIR/DebugImporter.h b/mlir/lib/Target/LLVMIR/DebugImporter.h
@@ -52,8 +52,15 @@ class DebugImporter {
   /// Infers the metadata type and translates it to MLIR.
   template <typename DINodeT>
   auto translate(DINodeT *node) {
-    // Infer the MLIR type from the LLVM metadata type.
-    using MLIRTypeT = decltype(translateImpl(node));
+    // Infer the result MLIR type from the LLVM metadata type.
+    // If the result is a DIType, it can also be wrapped in a recursive type,
+    // so the result is wrapped into a DIRecursiveTypeAttrOf.
+    // Otherwise, the exact result type is used.
+    constexpr bool isDIType = std::is_base_of_v<llvm::DIType, DINodeT>;
+    using RawMLIRTypeT = decltype(translateImpl(node));
+    using MLIRTypeT =
+        std::conditional_t<isDIType, DIRecursiveTypeAttrOf<RawMLIRTypeT>,
+                           RawMLIRTypeT>;
     return cast_or_null<MLIRTypeT>(
         translate(static_cast<llvm::DINode *>(node)));
   }
@@ -82,12 +89,27 @@ class DebugImporter {
   /// null attribute otherwise.
   StringAttr getStringAttrOrNull(llvm::MDString *stringNode);
 
+  DistinctAttr getOrCreateDistinctID(llvm::DINode *node);
+
   /// A mapping between LLVM debug metadata and the corresponding attribute.
   DenseMap<llvm::DINode *, DINodeAttr> nodeToAttr;
-
-  /// A stack that stores the metadata nodes that are being traversed. The stack
-  /// is used to detect cyclic dependencies during the metadata translation.
-  SetVector<llvm::DINode *> translationStack;
+  /// A mapping between LLVM debug metadata and the distinct ID attr for DI
+  /// nodes that require distinction.
+  DenseMap<llvm::DINode *, DistinctAttr> nodeToDistinctAttr;
+
+  /// A stack that stores the metadata type nodes that are being traversed. The
+  /// stack is used to detect cyclic dependencies during the metadata
+  /// translation. Nodes are pushed with a null value. If it is ever seen twice,
+  /// it is given a DistinctAttr, indicating that it is a recursive node and
+  /// should take on that DistinctAttr as ID.
+  llvm::MapVector<llvm::DIType *, DistinctAttr> typeTranslationStack;
+  /// All the unbound recursive self references in the translation stack.
+  SmallVector<DenseSet<DistinctAttr>> unboundRecursiveSelfRefs;
+  /// A stack that stores the non-type metadata nodes that are being traversed.
+  /// Each node is associated with the size of the `typeTranslationStack` at the
+  /// time of push. This is used to identify a recursion purely in the non-type
+  /// metadata nodes, which is not supported yet.
+  SetVector<std::pair<llvm::DINode *, unsigned>> nonTypeTranslationStack;
 
   MLIRContext *context;
   ModuleOp mlirModule;
diff --git a/mlir/test/Target/LLVMIR/Import/debug-info.ll b/mlir/test/Target/LLVMIR/Import/debug-info.ll
