[MLIR][LLVM] Add flag to skip import of DICompositeType's elems (#89355)

This commit introduces a flag to allow skipping the potentially
recursive import of DICompositeType elements. This patch is essentially a
bandaid for the still broken recursive debug type import.

Some of our downstream inputs are produced by excessive usage of
template meta programming, and thus contain tens of thousands of types
that all participate in such recursions. Unfortunately, the series of
patches that introduces type support is not easily revertible due to
being around for a while now and Modular depending on it.

We can consider to revert this change once the type importer has show to
be very performant, but for now we are talking second vs hours to import
specific files.

Author: Dinistro

mlir/include/mlir/Target/LLVMIR/Import.h

@@ -14,8 +14,6 @@
 #define MLIR_TARGET_LLVMIR_IMPORT_H
 
 #include "mlir/IR/OwningOpRef.h"
-#include "mlir/Support/LLVM.h"
-#include "llvm/ADT/StringRef.h"
 #include <memory>
 
 // Forward-declare LLVM classes.
@@ -34,12 +32,17 @@ class ModuleOp;
 /// The translation supports operations from any dialect that has a registered
 /// implementation of the LLVMImportDialectInterface. It returns nullptr if the
 /// translation fails and reports errors using the error handler registered with
-/// the MLIR context. The `emitExpensiveWarnings` option controls if expensive
+/// the MLIR context.
+/// The `emitExpensiveWarnings` option controls if expensive
 /// but uncritical diagnostics should be emitted.
+/// The `dropDICompositeTypeElements` option controls if DICompositeTypes should
+/// be imported without elements. If set, the option avoids the recursive
+/// traversal of composite type debug information, which can be expensive for
+/// adversarial inputs.
 OwningOpRef<ModuleOp>
 translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
-                        MLIRContext *context,
-                        bool emitExpensiveWarnings = true);
+                        MLIRContext *context, bool emitExpensiveWarnings = true,
+                        bool dropDICompositeTypeElements = false);
 
 /// Translate the given LLVM data layout into an MLIR equivalent using the DLTI
 /// dialect.

mlir/include/mlir/Target/LLVMIR/ModuleImport.h

@@ -47,7 +47,7 @@ class LoopAnnotationImporter;
 class ModuleImport {
 public:
   ModuleImport(ModuleOp mlirModule, std::unique_ptr<llvm::Module> llvmModule,
-               bool emitExpensiveWarnings);
+               bool emitExpensiveWarnings, bool importEmptyDICompositeTypes);
 
   /// Calls the LLVMImportInterface initialization that queries the registered
   /// dialect interfaces for the supported LLVM IR intrinsics and metadata kinds

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

@@ -30,6 +30,12 @@ void registerFromLLVMIRTranslation() {
       llvm::cl::desc("Emit expensive warnings during LLVM IR import "
                      "(discouraged: testing only!)"),
       llvm::cl::init(false));
+  static llvm::cl::opt<bool> dropDICompositeTypeElements(
+      "drop-di-composite-type-elements",
+      llvm::cl::desc(
+          "Avoid translating the elements of DICompositeTypes during "
+          "the LLVM IR import (discouraged: testing only!)"),
+      llvm::cl::init(false));
 
   TranslateToMLIRRegistration registration(
       "import-llvm", "Translate LLVMIR to MLIR",
@@ -51,7 +57,8 @@ void registerFromLLVMIRTranslation() {
           return nullptr;
 
         return translateLLVMIRToModule(std::move(llvmModule), context,
-                                       emitExpensiveWarnings);
+                                       emitExpensiveWarnings,
+                                       dropDICompositeTypeElements);
       },
       [](DialectRegistry &registry) {
         // Register the DLTI dialect used to express the data layout

mlir/lib/Target/LLVMIR/DebugImporter.cpp

@@ -26,9 +26,11 @@ using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
 
-DebugImporter::DebugImporter(ModuleOp mlirModule)
+DebugImporter::DebugImporter(ModuleOp mlirModule,
+                             bool dropDICompositeTypeElements)
     : recursionPruner(mlirModule.getContext()),
-      context(mlirModule.getContext()), mlirModule(mlirModule) {}
+      context(mlirModule.getContext()), mlirModule(mlirModule),
+      dropDICompositeTypeElements(dropDICompositeTypeElements) {}
 
 Location DebugImporter::translateFuncLocation(llvm::Function *func) {
   llvm::DISubprogram *subprogram = func->getSubprogram();
@@ -69,9 +71,11 @@ DICompileUnitAttr DebugImporter::translateImpl(llvm::DICompileUnit *node) {
 DICompositeTypeAttr DebugImporter::translateImpl(llvm::DICompositeType *node) {
   std::optional<DIFlags> flags = symbolizeDIFlags(node->getFlags());
   SmallVector<DINodeAttr> elements;
-  for (llvm::DINode *element : node->getElements()) {
-    assert(element && "expected a non-null element type");
-    elements.push_back(translate(element));
+  if (!dropDICompositeTypeElements) {
+    for (llvm::DINode *element : node->getElements()) {
+      assert(element && "expected a non-null element type");
+      elements.push_back(translate(element));
+    }
   }
   // Drop the elements parameter if any of the elements are invalid.
   if (llvm::is_contained(elements, nullptr))

mlir/lib/Target/LLVMIR/DebugImporter.h

@@ -29,7 +29,7 @@ namespace detail {
 
 class DebugImporter {
 public:
-  DebugImporter(ModuleOp mlirModule);
+  DebugImporter(ModuleOp mlirModule, bool dropDICompositeTypeElements);
 
   /// Translates the given LLVM debug location to an MLIR location.
   Location translateLoc(llvm::DILocation *loc);
@@ -184,6 +184,12 @@ class DebugImporter {
 
   MLIRContext *context;
   ModuleOp mlirModule;
+
+  /// An option to control if DICompositeTypes should always be imported without
+  /// converting their elements. If set, the option avoids the recursive
+  /// traversal of composite type debug information, which can be expensive for
+  /// adversarial inputs.
+  bool dropDICompositeTypeElements;
 };
 
 } // namespace detail

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -155,12 +155,14 @@ getTopologicallySortedBlocks(ArrayRef<llvm::BasicBlock *> basicBlocks) {
 
 ModuleImport::ModuleImport(ModuleOp mlirModule,
                            std::unique_ptr<llvm::Module> llvmModule,
-                           bool emitExpensiveWarnings)
+                           bool emitExpensiveWarnings,
+                           bool importEmptyDICompositeTypes)
     : builder(mlirModule->getContext()), context(mlirModule->getContext()),
       mlirModule(mlirModule), llvmModule(std::move(llvmModule)),
       iface(mlirModule->getContext()),
       typeTranslator(*mlirModule->getContext()),
-      debugImporter(std::make_unique<DebugImporter>(mlirModule)),
+      debugImporter(std::make_unique<DebugImporter>(
+          mlirModule, importEmptyDICompositeTypes)),
       loopAnnotationImporter(
           std::make_unique<LoopAnnotationImporter>(*this, builder)),
       emitExpensiveWarnings(emitExpensiveWarnings) {
@@ -2092,8 +2094,8 @@ ModuleImport::translateLoopAnnotationAttr(const llvm::MDNode *node,
 
 OwningOpRef<ModuleOp>
 mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
-                              MLIRContext *context,
-                              bool emitExpensiveWarnings) {
+                              MLIRContext *context, bool emitExpensiveWarnings,
+                              bool dropDICompositeTypeElements) {
   // Preload all registered dialects to allow the import to iterate the
   // registered LLVMImportDialectInterface implementations and query the
   // supported LLVM IR constructs before starting the translation. Assumes the
@@ -2109,7 +2111,7 @@ mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
       /*column=*/0)));
 
   ModuleImport moduleImport(module.get(), std::move(llvmModule),
-                            emitExpensiveWarnings);
+                            emitExpensiveWarnings, dropDICompositeTypeElements);
   if (failed(moduleImport.initializeImportInterface()))
     return {};
   if (failed(moduleImport.convertDataLayout()))

mlir/test/Target/LLVMIR/Import/ignore-composite-type-elements.ll

@@ -0,0 +1,25 @@
+; RUN: mlir-translate -import-llvm -mlir-print-debuginfo -split-input-file -drop-di-composite-type-elements %s | FileCheck %s
+
+; Verifies that the according flag avoids the conversion of the elements of the
+; DICompositeType.
+
+; CHECK-NOT: di_derive_type
+; CHECK: #{{.+}} = #llvm.di_composite_type<tag = DW_TAG_class_type, name = "class">
+; CHECK-NOT: di_derive_type
+
+define void @composite_type() !dbg !3 {
+  ret void
+}
+
+!llvm.dbg.cu = !{!1}
+!llvm.module.flags = !{!0}
+!0 = !{i32 2, !"Debug Info Version", i32 3}
+!1 = distinct !DICompileUnit(language: DW_LANG_C, file: !2)
+!2 = !DIFile(filename: "debug-info.ll", directory: "/")
+!3 = distinct !DISubprogram(name: "composite_type", scope: !2, file: !2, spFlags: DISPFlagDefinition, unit: !1, type: !4)
+!4 = !DISubroutineType(types: !5)
+!5 = !{!6}
+!6 = !DICompositeType(tag: DW_TAG_class_type, name: "class", elements: !7)
+!7 = !{!9}
+!8 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !6, flags: DIFlagArtificial | DIFlagObjectPointer)
+!9 = !DIDerivedType(tag: DW_TAG_member, name: "call_field", file: !2, baseType: !8)