Reapply [MLIR][LLVMIR] Import unregistered intrinsics via llvm.intrin… (#129174)

…sic_call

Original introduced in #128626,
reverted in #128973

Reproduced the issue on a shared lib build locally on Linux, moved
content around to satisfy both static and shared lib builds.

### Original commit message

Currently, the llvm importer can only cover intrinsics that have a first
class representation in an MLIR dialect (arm-neon, etc). This PR
introduces a fallback mechanism that allow "unregistered" intrinsics to
be imported by using the generic `llvm.intrinsic_call` operation. This
is useful in several ways:

1. Allows round-trip the LLVM dialect output lowered from other dialects
(example: ClangIR)
2. Enables MLIR-linking tools to operate on imported LLVM IR without
requiring to add new operations to dozen of different targets.

If multiple dialects implement this interface hook, the last one to
register is the one converting all unregistered intrinsics.

---------

Co-authored-by: Bruno Cardoso Lopes <[email protected]>

Author: bcardosolopes

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

@@ -155,9 +155,18 @@ class LLVMImportInterface
   LogicalResult convertIntrinsic(OpBuilder &builder, llvm::CallInst *inst,
                                  LLVM::ModuleImport &moduleImport) const {
     // Lookup the dialect interface for the given intrinsic.
-    Dialect *dialect = intrinsicToDialect.lookup(inst->getIntrinsicID());
+    // 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 failure();
+      return convertUnregisteredIntrinsic(builder, inst, moduleImport);
 
     // Dispatch the conversion to the dialect interface.
     const LLVMImportDialectInterface *iface = getInterfaceFor(dialect);
@@ -224,6 +233,11 @@ class LLVMImportInterface
   }
 
 private:
+  /// Generate llvm.call_intrinsic when no supporting dialect available.
+  static LogicalResult
+  convertUnregisteredIntrinsic(OpBuilder &builder, llvm::CallInst *inst,
+                               LLVM::ModuleImport &moduleImport);
+
   DenseMap<unsigned, Dialect *> intrinsicToDialect;
   DenseMap<unsigned, const LLVMImportDialectInterface *> instructionToDialect;
   DenseMap<unsigned, SmallVector<Dialect *, 1>> metadataToDialect;

mlir/lib/Target/LLVMIR/CMakeLists.txt

@@ -8,6 +8,7 @@ set(LLVM_OPTIONAL_SOURCES
   DebugImporter.cpp
   LoopAnnotationImporter.cpp
   LoopAnnotationTranslation.cpp
+  LLVMImportInterface.cpp
   ModuleTranslation.cpp
   ModuleImport.cpp
   TypeToLLVM.cpp
@@ -68,6 +69,7 @@ add_mlir_translation_library(MLIRTargetLLVMIRImport
   DebugImporter.cpp
   LoopAnnotationImporter.cpp
   ModuleImport.cpp
+  LLVMImportInterface.cpp
   TypeFromLLVM.cpp
 
   ADDITIONAL_HEADER_DIRS

mlir/lib/Target/LLVMIR/LLVMImportInterface.cpp

@@ -0,0 +1,59 @@
+//===------------------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements methods from LLVMImportInterface.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVMIR/LLVMImportInterface.h"
+#include "mlir/Target/LLVMIR/Import.h"
+#include "mlir/Target/LLVMIR/ModuleImport.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+using namespace mlir::LLVM::detail;
+
+LogicalResult mlir::LLVMImportInterface::convertUnregisteredIntrinsic(
+    OpBuilder &builder, llvm::CallInst *inst,
+    LLVM::ModuleImport &moduleImport) {
+  StringRef intrinName = inst->getCalledFunction()->getName();
+
+  SmallVector<llvm::Value *> args(inst->args());
+  ArrayRef<llvm::Value *> llvmOperands(args);
+
+  SmallVector<llvm::OperandBundleUse> llvmOpBundles;
+  llvmOpBundles.reserve(inst->getNumOperandBundles());
+  for (unsigned i = 0; i < inst->getNumOperandBundles(); ++i)
+    llvmOpBundles.push_back(inst->getOperandBundleAt(i));
+
+  SmallVector<Value> mlirOperands;
+  SmallVector<NamedAttribute> mlirAttrs;
+  if (failed(moduleImport.convertIntrinsicArguments(
+          llvmOperands, llvmOpBundles, false, {}, {}, mlirOperands, mlirAttrs)))
+    return failure();
+
+  Type results = moduleImport.convertType(inst->getType());
+  auto op = builder.create<::mlir::LLVM::CallIntrinsicOp>(
+      moduleImport.translateLoc(inst->getDebugLoc()), results,
+      StringAttr::get(builder.getContext(), intrinName),
+      ValueRange{mlirOperands}, FastmathFlagsAttr{});
+
+  moduleImport.setFastmathFlagsAttr(inst, op);
+
+  // Update importer tracking of results.
+  unsigned numRes = op.getNumResults();
+  if (numRes == 1)
+    moduleImport.mapValue(inst) = op.getResult(0);
+  else if (numRes == 0)
+    moduleImport.mapNoResultOp(inst);
+  else
+    return op.emitError(
+        "expected at most one result from target intrinsic call");
+
+  return success();
+}

mlir/test/Target/LLVMIR/Import/import-failure.ll

@@ -38,18 +38,6 @@ bb1:
 
 ; // -----
 
-declare void @llvm.gcroot(ptr %arg1, ptr %arg2)
-
-; CHECK:      <unknown>
-; CHECK-SAME: error: unhandled intrinsic: call void @llvm.gcroot(ptr %arg1, ptr null)
-define void @unhandled_intrinsic() gc "example" {
-  %arg1 = alloca ptr
-  call void @llvm.gcroot(ptr %arg1, ptr null)
-  ret void
-}
-
-; // -----
-
 ; Check that debug intrinsics with an unsupported argument are dropped.
 
 declare void @llvm.dbg.value(metadata, metadata, metadata)

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

@@ -0,0 +1,68 @@
+; RUN: mlir-translate -import-llvm %s -split-input-file | FileCheck %s
+
+declare i64 @llvm.aarch64.ldxr.p0(ptr)
+
+define dso_local void @t0(ptr %a) {
+  %x = call i64 @llvm.aarch64.ldxr.p0(ptr elementtype(i8) %a)
+  ret void
+}
+
+; CHECK-LABEL: llvm.func @llvm.aarch64.ldxr.p0(!llvm.ptr)
+; CHECK-LABEL: llvm.func @t0
+; CHECK:   llvm.call_intrinsic "llvm.aarch64.ldxr.p0"({{.*}}) : (!llvm.ptr) -> i64
+; CHECK:   llvm.return
+; CHECK: }
+
+; // -----
+
+declare <8 x i8> @llvm.aarch64.neon.uabd.v8i8(<8 x i8>, <8 x i8>)
+
+define dso_local <8 x i8> @t1(<8 x i8> %lhs, <8 x i8> %rhs) {
+  %r = call <8 x i8> @llvm.aarch64.neon.uabd.v8i8(<8 x i8> %lhs, <8 x i8> %rhs)
+  ret <8 x i8> %r
+}
+
+; CHECK: llvm.func @t1(%[[A0:.*]]: vector<8xi8>, %[[A1:.*]]: vector<8xi8>) -> vector<8xi8> {{.*}} {
+; CHECK:   %[[R:.*]] = llvm.call_intrinsic "llvm.aarch64.neon.uabd.v8i8"(%[[A0]], %[[A1]]) : (vector<8xi8>, vector<8xi8>) -> vector<8xi8>
+; CHECK:   llvm.return %[[R]] : vector<8xi8>
+; CHECK: }
+
+; // -----
+
+declare void @llvm.aarch64.neon.st2.v8i8.p0(<8 x i8>, <8 x i8>, ptr)
+
+define dso_local void @t2(<8 x i8> %lhs, <8 x i8> %rhs, ptr %a) {
+  call void @llvm.aarch64.neon.st2.v8i8.p0(<8 x i8> %lhs, <8 x i8> %rhs, ptr %a)
+  ret void
+}
+
+; CHECK: llvm.func @t2(%[[A0:.*]]: vector<8xi8>, %[[A1:.*]]: vector<8xi8>, %[[A2:.*]]: !llvm.ptr) {{.*}} {
+; CHECK:   llvm.call_intrinsic "llvm.aarch64.neon.st2.v8i8.p0"(%[[A0]], %[[A1]], %[[A2]]) : (vector<8xi8>, vector<8xi8>, !llvm.ptr) -> !llvm.void
+; CHECK:   llvm.return
+; CHECK: }
+
+; // -----
+
+declare void @llvm.gcroot(ptr %arg1, ptr %arg2)
+define void @gctest() gc "example" {
+  %arg1 = alloca ptr
+  call void @llvm.gcroot(ptr %arg1, ptr null)
+  ret void
+}
+
+; CHECK-LABEL: @gctest
+; CHECK: llvm.call_intrinsic "llvm.gcroot"({{.*}}, {{.*}}) : (!llvm.ptr, !llvm.ptr) -> !llvm.void
+
+; // -----
+
+; Test we get the supported version, not the unregistered one.
+
+declare i32 @llvm.lround.i32.f32(float)
+
+; CHECK-LABEL: llvm.func @lround_test
+define void @lround_test(float %0, double %1) {
+  ; CHECK-NOT: llvm.call_intrinsic "llvm.lround
+  ; CHECK: llvm.intr.lround(%{{.*}}) : (f32) -> i32
+  %3 = call i32 @llvm.lround.i32.f32(float %0)
+  ret void
+}