[MLIR][LLVM] Add import flag to skip traversal of DICompositeType's elems

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)