[MLIR][LLVMIR] Import: add flag to prefer using unregistered intrinsics

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.

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,13 @@ 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
+/// prefer generic version of imported intrinsics with `llvm.intrinsic_call`
+/// than using versions supported by a dialects.
+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
@@ -277,6 +278,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() {
@@ -474,6 +481,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 to disable supported intrinsic support in
+  /// favor a more generic version via `llvm.intrinsic_call`.
+  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 translation 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());
 }
 
@@ -2527,11 +2529,10 @@ ModuleImport::translateLoopAnnotationAttr(const llvm::MDNode *node,
   return loopAnnotationImporter->translateLoopAnnotation(node, loc);
 }
 
-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
@@ -2548,7 +2549,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
+}