[MLIR][LLVMIR] Import: add flag to prefer using unregistered intrinsics (#130685)

Currently, there is no common mechanism for supported intrinsics to be
generically annotated with arg and ret attributes. Since there are many
supported intrinsics around different dialects, the amount of work to
teach all them about these attributes is not trivial (though it would be
nice in the long term).

This PR adds a new flag `-prefer-unregistered-intrinsics` that can be
used alongside `--import-llvm` to always use `llvm.intrinsic_call`
during import time (ignoring dialect hooks for custom intrinsic
support).

Using this flag allow us to roundtrip the LLVM IR while eliminating a
whole set of differences coming from lack of arg/ret attributes on
supported intrinsics.

Note `convertIntrinsic` has to be moved to an implementation file
because it queries into `moduleImport` state, which is a fwd declaration
in `LLVMImportInterface.h`

Author: bcardosolopes

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

@@ -41,11 +41,15 @@ class ModuleOp;
 /// adversarial inputs.
 /// The `loadAllDialects` flag (default on) will load all dialects in the
 /// context.
-OwningOpRef<ModuleOp>
-translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
-                        MLIRContext *context, bool emitExpensiveWarnings = true,
-                        bool dropDICompositeTypeElements = false,
-                        bool loadAllDialects = true);
+/// The `preferUnregisteredIntrinsics` flag (default off) controls whether to
+/// import all intrinsics using `llvm.intrinsic_call` even if a dialect
+/// registered an explicit intrinsic operation. Warning: passes that rely on
+/// matching explicit intrinsic operations may not work properly if this flag is
+/// enabled.
+OwningOpRef<ModuleOp> translateLLVMIRToModule(
+    std::unique_ptr<llvm::Module> llvmModule, MLIRContext *context,
+    bool emitExpensiveWarnings = true, bool dropDICompositeTypeElements = false,
+    bool loadAllDialects = true, bool preferUnregisteredIntrinsics = false);
 
 /// Translate the given LLVM data layout into an MLIR equivalent using the DLTI
 /// dialect.

mlir/include/mlir/Target/LLVMIR/LLVMImportInterface.h

@@ -153,26 +153,7 @@ class LLVMImportInterface
   /// Converts the LLVM intrinsic to an MLIR operation if a conversion exists.
   /// Returns failure otherwise.
   LogicalResult convertIntrinsic(OpBuilder &builder, llvm::CallInst *inst,
-                                 LLVM::ModuleImport &moduleImport) const {
-    // Lookup the dialect interface for the given intrinsic.
-    // Verify the intrinsic identifier maps to an actual intrinsic.
-    llvm::Intrinsic::ID intrinId = inst->getIntrinsicID();
-    assert(intrinId != llvm::Intrinsic::not_intrinsic);
-
-    // First lookup the intrinsic across different dialects for known
-    // supported conversions, examples include arm-neon, nvm-sve, etc.
-    Dialect *dialect = intrinsicToDialect.lookup(intrinId);
-
-    // No specialized (supported) intrinsics, attempt to generate a generic
-    // version via llvm.call_intrinsic (if available).
-    if (!dialect)
-      return convertUnregisteredIntrinsic(builder, inst, moduleImport);
-
-    // Dispatch the conversion to the dialect interface.
-    const LLVMImportDialectInterface *iface = getInterfaceFor(dialect);
-    assert(iface && "expected to find a dialect interface");
-    return iface->convertIntrinsic(builder, inst, moduleImport);
-  }
+                                 LLVM::ModuleImport &moduleImport) const;
 
   /// Returns true if the given LLVM IR intrinsic is convertible to an MLIR
   /// operation.

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

@@ -47,7 +47,8 @@ class LoopAnnotationImporter;
 class ModuleImport {
 public:
   ModuleImport(ModuleOp mlirModule, std::unique_ptr<llvm::Module> llvmModule,
-               bool emitExpensiveWarnings, bool importEmptyDICompositeTypes);
+               bool emitExpensiveWarnings, bool importEmptyDICompositeTypes,
+               bool preferUnregisteredIntrinsics);
 
   /// Calls the LLVMImportInterface initialization that queries the registered
   /// dialect interfaces for the supported LLVM IR intrinsics and metadata kinds
@@ -287,6 +288,12 @@ class ModuleImport {
   void convertParameterAttributes(llvm::CallBase *call, ArrayAttr &argsAttr,
                                   ArrayAttr &resAttr, OpBuilder &builder);
 
+  /// Whether the importer should try to convert all intrinsics to
+  /// llvm.call_intrinsic instead of dialect supported operations.
+  bool useUnregisteredIntrinsicsOnly() const {
+    return preferUnregisteredIntrinsics;
+  }
+
 private:
   /// Clears the accumulated state before processing a new region.
   void clearRegionState() {
@@ -484,6 +491,10 @@ class ModuleImport {
   /// emitted. Avoids generating warnings for unhandled debug intrinsics and
   /// metadata that otherwise dominate the translation time for large inputs.
   bool emitExpensiveWarnings;
+
+  /// An option to control whether the importer should try to convert all
+  /// intrinsics to llvm.call_intrinsic instead of dialect supported operations.
+  bool preferUnregisteredIntrinsics;
 };
 
 } // namespace LLVM

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

@@ -37,6 +37,13 @@ void registerFromLLVMIRTranslation() {
           "the LLVM IR import (discouraged: testing only!)"),
       llvm::cl::init(false));
 
+  static llvm::cl::opt<bool> preferUnregisteredIntrinsics(
+      "prefer-unregistered-intrinsics",
+      llvm::cl::desc(
+          "Prefer translating all intrinsics into llvm.call_intrinsic instead "
+          "of using dialect supported intrinsics"),
+      llvm::cl::init(false));
+
   TranslateToMLIRRegistration registration(
       "import-llvm", "Translate LLVMIR to MLIR",
       [](llvm::SourceMgr &sourceMgr,
@@ -60,9 +67,10 @@ void registerFromLLVMIRTranslation() {
         if (llvmModule->IsNewDbgInfoFormat)
           llvmModule->convertFromNewDbgValues();
 
-        return translateLLVMIRToModule(std::move(llvmModule), context,
-                                       emitExpensiveWarnings,
-                                       dropDICompositeTypeElements);
+        return translateLLVMIRToModule(
+            std::move(llvmModule), context, emitExpensiveWarnings,
+            dropDICompositeTypeElements, /*loadAllDialects=*/true,
+            preferUnregisteredIntrinsics);
       },
       [](DialectRegistry &registry) {
         // Register the DLTI dialect used to express the data layout

mlir/lib/Target/LLVMIR/LLVMImportInterface.cpp

@@ -62,3 +62,31 @@ LogicalResult mlir::LLVMImportInterface::convertUnregisteredIntrinsic(
 
   return success();
 }
+
+/// Converts the LLVM intrinsic to an MLIR operation if a conversion exists.
+/// Returns failure otherwise.
+LogicalResult mlir::LLVMImportInterface::convertIntrinsic(
+    OpBuilder &builder, llvm::CallInst *inst,
+    LLVM::ModuleImport &moduleImport) const {
+  // Lookup the dialect interface for the given intrinsic.
+  // Verify the intrinsic identifier maps to an actual intrinsic.
+  llvm::Intrinsic::ID intrinId = inst->getIntrinsicID();
+  assert(intrinId != llvm::Intrinsic::not_intrinsic);
+
+  // First lookup the intrinsic across different dialects for known
+  // supported conversions, examples include arm-neon, nvm-sve, etc.
+  Dialect *dialect = nullptr;
+
+  if (!moduleImport.useUnregisteredIntrinsicsOnly())
+    dialect = intrinsicToDialect.lookup(intrinId);
+
+  // No specialized (supported) intrinsics, attempt to generate a generic
+  // version via llvm.call_intrinsic (if available).
+  if (!dialect)
+    return convertUnregisteredIntrinsic(builder, inst, moduleImport);
+
+  // Dispatch the conversion to the dialect interface.
+  const LLVMImportDialectInterface *iface = getInterfaceFor(dialect);
+  assert(iface && "expected to find a dialect interface");
+  return iface->convertIntrinsic(builder, inst, moduleImport);
+}

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -162,7 +162,8 @@ getTopologicallySortedBlocks(ArrayRef<llvm::BasicBlock *> basicBlocks) {
 ModuleImport::ModuleImport(ModuleOp mlirModule,
                            std::unique_ptr<llvm::Module> llvmModule,
                            bool emitExpensiveWarnings,
-                           bool importEmptyDICompositeTypes)
+                           bool importEmptyDICompositeTypes,
+                           bool preferUnregisteredIntrinsics)
     : builder(mlirModule->getContext()), context(mlirModule->getContext()),
       mlirModule(mlirModule), llvmModule(std::move(llvmModule)),
       iface(mlirModule->getContext()),
@@ -171,7 +172,8 @@ ModuleImport::ModuleImport(ModuleOp mlirModule,
           mlirModule, importEmptyDICompositeTypes)),
       loopAnnotationImporter(
           std::make_unique<LoopAnnotationImporter>(*this, builder)),
-      emitExpensiveWarnings(emitExpensiveWarnings) {
+      emitExpensiveWarnings(emitExpensiveWarnings),
+      preferUnregisteredIntrinsics(preferUnregisteredIntrinsics) {
   builder.setInsertionPointToStart(mlirModule.getBody());
 }
 
@@ -2581,11 +2583,10 @@ ModuleImport::translateDereferenceableAttr(const llvm::MDNode *node,
   return derefAttr;
 }
 
-OwningOpRef<ModuleOp>
-mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
-                              MLIRContext *context, bool emitExpensiveWarnings,
-                              bool dropDICompositeTypeElements,
-                              bool loadAllDialects) {
+OwningOpRef<ModuleOp> mlir::translateLLVMIRToModule(
+    std::unique_ptr<llvm::Module> llvmModule, MLIRContext *context,
+    bool emitExpensiveWarnings, bool dropDICompositeTypeElements,
+    bool loadAllDialects, bool preferUnregisteredIntrinsics) {
   // 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
@@ -2602,7 +2603,8 @@ mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
       /*column=*/0)));
 
   ModuleImport moduleImport(module.get(), std::move(llvmModule),
-                            emitExpensiveWarnings, dropDICompositeTypeElements);
+                            emitExpensiveWarnings, dropDICompositeTypeElements,
+                            preferUnregisteredIntrinsics);
   if (failed(moduleImport.initializeImportInterface()))
     return {};
   if (failed(moduleImport.convertDataLayout()))

mlir/test/Target/LLVMIR/Import/intrinsic-prefer-unregistered.ll

@@ -0,0 +1,10 @@
+; RUN: mlir-translate -import-llvm -prefer-unregistered-intrinsics %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @lifetime
+define void @lifetime(ptr %0) {
+  ; CHECK: llvm.call_intrinsic "llvm.lifetime.start.p0"({{.*}}, %arg0) : (i64, !llvm.ptr {llvm.nonnull}) -> !llvm.void
+  call void @llvm.lifetime.start.p0(i64 16, ptr nonnull %0)
+  ; CHECK: llvm.call_intrinsic "llvm.lifetime.end.p0"({{.*}}, %arg0) : (i64, !llvm.ptr {llvm.nonnull}) -> !llvm.void
+  call void @llvm.lifetime.end.p0(i64 32, ptr nonnull %0)
+  ret void
+}