[MLIR][LLVMIR] Import unregistered intrinsics via llvm.intrinsic_call (#128626)

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 (cc
@xlauko @smeenai).

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

---------

Co-authored-by: Tobias Gysi <[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/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

@@ -367,6 +367,45 @@ static LogicalResult setIntelReqdSubGroupSizeAttr(Builder &builder,
   return success();
 }
 
+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();
+}
 namespace {
 
 /// Implementation of the dialect interface that converts operations belonging

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
+}