[mlir][llvmir] Import intrinsics with attributes from LLVMIR.

The revision adds support to specify custom import functions for
LLVM IR intrinsics with immediate arguments that translate to MLIR
attributes. It takes an approach similar to the MLIR to LLVM translation
that uses a tablegen defined build method. The default implementation
of this newly introduced "mlirBuilder" assumes all intrinsic arguments
translate to operands. Specific intrinsics, such as
llvm.lifetime.start/stop then define a custom builder that converts
their immediate arguments to MLIR attributes.

Depends on D135349

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D135350

Author: gysit

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

@@ -37,7 +37,8 @@ add_public_tablegen_target(MLIRLLVMConversionsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LLVMIntrinsicOps.td)
 mlir_tablegen(LLVMIntrinsicConversions.inc -gen-llvmir-conversions)
-mlir_tablegen(LLVMIntrinsicToLLVMIROpPairs.inc -gen-llvmintrinsic-to-llvmirop-pairs)
+mlir_tablegen(LLVMIntrinsicFromLLVMIRConversions.inc -gen-intr-from-llvmir-conversions)
+mlir_tablegen(LLVMConvertibleLLVMIRIntrinsics.inc -gen-convertible-llvmir-intrinsics)
 add_public_tablegen_target(MLIRLLVMIntrinsicConversionsIncGen)
 
 add_mlir_dialect(NVVMOps nvvm)
@@ -55,4 +56,3 @@ add_mlir_doc(ROCDLOps ROCDLDialect Dialects/ -gen-dialect-doc -dialect=rocdl)
 set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
 mlir_tablegen(ROCDLConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRROCDLConversionsIncGen)
-

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

@@ -121,8 +121,20 @@ class LLVM_LifetimeBaseOp<string opName> : LLVM_ZeroResultIntrOp<opName> {
   let assemblyFormat = "$size `,` $ptr attr-dict `:` type($ptr)";
 }
 
-def LLVM_LifetimeStartOp : LLVM_LifetimeBaseOp<"lifetime.start">;
-def LLVM_LifetimeEndOp : LLVM_LifetimeBaseOp<"lifetime.end">;
+def LLVM_LifetimeStartOp : LLVM_LifetimeBaseOp<"lifetime.start"> {
+  // Custom builder to convert the size argument to an attribute.
+  string mlirBuilder = [{
+    $_builder.create<LLVM::LifetimeStartOp>(
+      $_location, $_int_attr($size), $ptr);
+  }];
+}
+def LLVM_LifetimeEndOp : LLVM_LifetimeBaseOp<"lifetime.end"> {
+  // Custom builder to convert the size argument to an attribute.
+  string mlirBuilder = [{
+    $_builder.create<LLVM::LifetimeEndOp>(
+      $_location, $_int_attr($size), $ptr);
+  }];
+}
 
 // Intrinsics with multiple returns.
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

@@ -302,7 +302,9 @@ def LLVM_IntrPatterns {
 // to be accessed via the LLVMStructType, instead of directly via the result.
 // "opName" contains the name of the operation to be associated with the
 // intrinsic and "enumName" contains the name of the intrinsic as appears in
-// `llvm::Intrinsic` enum; one usually wants these to be related.
+// `llvm::Intrinsic` enum; one usually wants these to be related. Additionally,
+// the base class also defines the "mlirBuilder" field to support the inverse
+// translation starting from an LLVM IR intrinsic.
 class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
                       list<int> overloadedResults, list<int> overloadedOperands,
                       list<Trait> traits, int numResults,
@@ -312,7 +314,7 @@ class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
   string resultPattern = !if(!gt(numResults, 1),
                              LLVM_IntrPatterns.structResult,
                              LLVM_IntrPatterns.result);
-  string llvmEnumName = enumName;                           
+  string llvmEnumName = enumName;
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
@@ -332,6 +334,30 @@ class LLVM_IntrOpBase<Dialect dialect, string opName, string enumName,
       "moduleTranslation.setAliasScopeMetadata(op, inst);",
       "(void) inst;")
     # !if(!gt(numResults, 0), "$res = inst;", "");
+
+  // A builder to construct the MLIR LLVM dialect operation given the matching
+  // LLVM IR intrinsic instruction `callInst`. The following $-variables exist:
+  //   - $name - substituted by the remapped `callInst` argument using the
+  //             the index of the MLIR operation argument with the given name;
+  //   - $_int_attr - substituted by a call to an integer attribute matcher;
+  //   - $_resultType - substituted with the MLIR result type;
+  //   - $_location - substituted with the MLIR location;
+  //   - $_builder - substituted with the MLIR builder;
+  //   - $_qualCppClassName - substitiuted with the MLIR operation class name.
+  // Additionally, `$$` can be used to produce the dollar character.
+  // NOTE: The $name variable resolution assumes the MLIR and LLVM argument
+  // orders match and there are no optional or variadic arguments.
+  string mlirBuilder = [{
+    SmallVector<llvm::Value *> operands(callInst->args());
+    SmallVector<Type> resultTypes =
+    }] # !if(!gt(numResults, 0),
+      "{$_resultType};", "{};") # [{
+    Operation *op = $_builder.create<$_qualCppClassName>(
+      $_location,
+      resultTypes,
+      processValues(operands));
+    }] # !if(!gt(numResults, 0),
+      "instMap[callInst] = op->getResult(0);", "(void)op;");
 }
 
 // Base class for LLVM intrinsic operations, should not be used directly. Places

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

@@ -18,6 +18,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Target/LLVMIR/TypeFromLLVM.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
@@ -31,6 +32,7 @@
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IRReader/IRReader.h"
@@ -42,6 +44,15 @@ using namespace mlir::LLVM;
 
 #include "mlir/Dialect/LLVMIR/LLVMConversionEnumsFromLLVM.inc"
 
+/// Returns true if the LLVM IR intrinsic is convertible to an MLIR LLVM dialect
+/// intrinsic, or false if no counterpart exists.
+static bool isConvertibleIntrinsic(llvm::Intrinsic::ID id) {
+  static const DenseSet<unsigned> convertibleIntrinsics = {
+#include "mlir/Dialect/LLVMIR/LLVMConvertibleLLVMIRIntrinsics.inc"
+  };
+  return convertibleIntrinsics.contains(id);
+}
+
 // Utility to print an LLVM value as a string for passing to emitError().
 // FIXME: Diagnostic should be able to natively handle types that have
 // operator << (raw_ostream&) defined.
@@ -173,13 +184,21 @@ class Importer {
   /// visited.
   SmallVector<Value> processValues(ArrayRef<llvm::Value *> values);
 
+  /// Converts `value` to an integer attribute. Asserts if the conversion fails.
+  IntegerAttr matchIntegerAttr(Value value);
+
   /// Translate the debug location to a FileLineColLoc, if `loc` is non-null.
   /// Otherwise, return UnknownLoc.
   Location translateLoc(llvm::DILocation *loc);
 
   /// Converts the type from LLVM to MLIR LLVM dialect.
   Type convertType(llvm::Type *type);
 
+  /// Converts an LLVM intrinsic to an MLIR LLVM dialect operation if an MLIR
+  /// counterpart exists. Otherwise, returns failure.
+  LogicalResult convertIntrinsic(OpBuilder &odsBuilder,
+                                 llvm::CallInst *callInst);
+
   /// Imports `f` into the current module.
   LogicalResult processFunction(llvm::Function *f);
 
@@ -263,6 +282,22 @@ Type Importer::convertType(llvm::Type *type) {
   return typeTranslator.translateType(type);
 }
 
+LogicalResult Importer::convertIntrinsic(OpBuilder &odsBuilder,
+                                         llvm::CallInst *callInst) {
+  llvm::Function *callee = callInst->getCalledFunction();
+  if (!callee || !callee->isIntrinsic())
+    return failure();
+
+  // Check if the intrinsic is convertible to an MLIR dialect counterpart.
+  llvm::Intrinsic::ID intrinsicID = callee->getIntrinsicID();
+  if (!isConvertibleIntrinsic(intrinsicID))
+    return failure();
+
+#include "mlir/Dialect/LLVMIR/LLVMIntrinsicFromLLVMIRConversions.inc"
+
+  return failure();
+}
+
 // We only need integers, floats, doubles, and vectors and tensors thereof for
 // attributes. Scalar and vector types are converted to the standard
 // equivalents. Array types are converted to ranked tensors; nested array types
@@ -568,6 +603,14 @@ SmallVector<Value> Importer::processValues(ArrayRef<llvm::Value *> values) {
   return remapped;
 }
 
+IntegerAttr Importer::matchIntegerAttr(Value value) {
+  IntegerAttr integerAttr;
+  bool success = matchPattern(value, m_Constant(&integerAttr));
+  assert(success && "expected a constant value");
+  (void)success;
+  return integerAttr;
+}
+
 /// Return the MLIR OperationName for the given LLVM opcode.
 static StringRef lookupOperationNameFromOpcode(unsigned opcode) {
 // Maps from LLVM opcode to MLIR OperationName. This is deliberately ordered
@@ -649,15 +692,6 @@ static StringRef lookupOperationNameFromOpcode(unsigned opcode) {
   return opcMap.lookup(opcode);
 }
 
-/// Return the MLIR OperationName for the given LLVM intrinsic ID.
-static StringRef lookupOperationNameFromIntrinsicID(unsigned id) {
-  // Maps from LLVM intrinsic ID to MLIR OperationName.
-  static const DenseMap<unsigned, StringRef> intrMap = {
-#include "mlir/Dialect/LLVMIR/LLVMIntrinsicToLLVMIROpPairs.inc"
-  };
-  return intrMap.lookup(id);
-}
-
 static ICmpPredicate getICmpPredicate(llvm::CmpInst::Predicate p) {
   switch (p) {
   default:
@@ -959,6 +993,11 @@ LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
   }
   case llvm::Instruction::Call: {
     llvm::CallInst *ci = cast<llvm::CallInst>(inst);
+
+    // For all intrinsics, try to generate to the corresponding op.
+    if (succeeded(convertIntrinsic(b, ci)))
+      return success();
+
     SmallVector<llvm::Value *> args(ci->args());
     SmallVector<Value> ops = processValues(args);
     SmallVector<Type, 2> tys;
@@ -968,20 +1007,6 @@ LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
     }
     Operation *op;
     if (llvm::Function *callee = ci->getCalledFunction()) {
-      // For all intrinsics, try to generate to the corresponding op.
-      if (callee->isIntrinsic()) {
-        auto id = callee->getIntrinsicID();
-        StringRef opName = lookupOperationNameFromIntrinsicID(id);
-        if (!opName.empty()) {
-          OperationState state(loc, opName);
-          state.addOperands(ops);
-          state.addTypes(tys);
-          Operation *op = b.create(state);
-          if (!inst->getType()->isVoidTy())
-            instMap[inst] = op->getResult(0);
-          return success();
-        }
-      }
       op = b.create<CallOp>(
           loc, tys, SymbolRefAttr::get(b.getContext(), callee->getName()), ops);
     } else {
@@ -1171,12 +1196,8 @@ LogicalResult Importer::processFunction(llvm::Function *f) {
 
   auto functionType =
       convertType(f->getFunctionType()).dyn_cast<LLVMFunctionType>();
-  if (f->isIntrinsic()) {
-    StringRef opName = lookupOperationNameFromIntrinsicID(f->getIntrinsicID());
-    // Skip the intrinsic decleration if we could found a corresponding op.
-    if (!opName.empty())
-      return success();
-  }
+  if (f->isIntrinsic() && isConvertibleIntrinsic(f->getIntrinsicID()))
+    return success();
 
   bool dsoLocal = f->hasLocalLinkage();
   CConv cconv = convertCConvFromLLVM(f->getCallingConv());

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

@@ -234,13 +234,13 @@ define void @umin_test(i32 %0, i32 %1, <8 x i32> %2, <8 x i32> %3) {
 define void @vector_reductions(float %0, <8 x float> %1, <8 x i32> %2) {
   ; CHECK: "llvm.intr.vector.reduce.add"(%{{.*}}) : (vector<8xi32>) -> i32
   %4 = call i32 @llvm.vector.reduce.add.v8i32(<8 x i32> %2)
-  ; CHECK: "llvm.intr.vector.reduce.and"(%{{.*}}) : (vector<8xi32>) -> i32  
+  ; CHECK: "llvm.intr.vector.reduce.and"(%{{.*}}) : (vector<8xi32>) -> i32
   %5 = call i32 @llvm.vector.reduce.and.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.fmax"(%{{.*}}) : (vector<8xf32>) -> f32
   %6 = call float @llvm.vector.reduce.fmax.v8f32(<8 x float> %1)
   ; CHECK: "llvm.intr.vector.reduce.fmin"(%{{.*}}) : (vector<8xf32>) -> f32
   %7 = call float @llvm.vector.reduce.fmin.v8f32(<8 x float> %1)
-  ; CHECK: "llvm.intr.vector.reduce.mul"(%{{.*}}) : (vector<8xi32>) -> i32  
+  ; CHECK: "llvm.intr.vector.reduce.mul"(%{{.*}}) : (vector<8xi32>) -> i32
   %8 = call i32 @llvm.vector.reduce.mul.v8i32(<8 x i32> %2)
   ; CHECK: "llvm.intr.vector.reduce.or"(%{{.*}}) : (vector<8xi32>) -> i32
   %9 = call i32 @llvm.vector.reduce.or.v8i32(<8 x i32> %2)
@@ -447,7 +447,7 @@ define void @coro_suspend(i32 %0, i1 %1, i8* %2) {
 
 ; CHECK-LABEL:  llvm.func @coro_end
 define void @coro_end(i8* %0, i1 %1) {
-  ; CHECK:  llvm.intr.coro.end 
+  ; CHECK:  llvm.intr.coro.end
   call i1 @llvm.coro.end(i8* %0, i1 %1)
   ret void
 }
@@ -489,11 +489,20 @@ define void @stack_restore(i8* %0) {
   ret void
 }
 
+; CHECK-LABEL: llvm.func @lifetime
+define void @lifetime(i8* %0) {
+  ; CHECK: llvm.intr.lifetime.start 16, %{{.*}} : !llvm.ptr<i8>
+  call void @llvm.lifetime.start.p0i8(i64 16, i8* %0)
+  ; CHECK: llvm.intr.lifetime.end 32, %{{.*}} : !llvm.ptr<i8>
+  call void @llvm.lifetime.end.p0i8(i64 32, i8* %0)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @vector_predication_intrinsics
 define void @vector_predication_intrinsics(<8 x i32> %0, <8 x i32> %1, <8 x float> %2, <8 x float> %3, <8 x i64> %4, <8 x double> %5, <8 x i32*> %6, i32 %7, float %8, i32* %9, float* %10, <8 x i1> %11, i32 %12) {
   ; CHECK: "llvm.intr.vp.add"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>
   %14 = call <8 x i32> @llvm.vp.add.v8i32(<8 x i32> %0, <8 x i32> %1, <8 x i1> %11, i32 %12)
-  ; CHECK: "llvm.intr.vp.sub"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>  
+  ; CHECK: "llvm.intr.vp.sub"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>
   %15 = call <8 x i32> @llvm.vp.sub.v8i32(<8 x i32> %0, <8 x i32> %1, <8 x i1> %11, i32 %12)
   ; CHECK: "llvm.intr.vp.mul"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi32>, vector<8xi32>, vector<8xi1>, i32) -> vector<8xi32>
   %16 = call <8 x i32> @llvm.vp.mul.v8i32(<8 x i32> %0, <8 x i32> %1, <8 x i1> %11, i32 %12)
@@ -585,7 +594,7 @@ define void @vector_predication_intrinsics(<8 x i32> %0, <8 x i32> %1, <8 x floa
   %57 = call <8 x i64> @llvm.vp.fptosi.v8i64.v8f64(<8 x double> %5, <8 x i1> %11, i32 %12)
   ; CHECK: "llvm.intr.vp.ptrtoint"(%{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.vec<8 x ptr<i32>>, vector<8xi1>, i32) -> vector<8xi64>
   %58 = call <8 x i64> @llvm.vp.ptrtoint.v8i64.v8p0i32(<8 x i32*> %6, <8 x i1> %11, i32 %12)
-  ; CHECK: "llvm.intr.vp.inttoptr"(%{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi64>, vector<8xi1>, i32) -> !llvm.vec<8 x ptr<i32>>  
+  ; CHECK: "llvm.intr.vp.inttoptr"(%{{.*}}, %{{.*}}, %{{.*}}) : (vector<8xi64>, vector<8xi1>, i32) -> !llvm.vec<8 x ptr<i32>>
   %59 = call <8 x i32*> @llvm.vp.inttoptr.v8p0i32.v8i64(<8 x i64> %4, <8 x i1> %11, i32 %12)
   ret void
 }
@@ -662,7 +671,7 @@ declare <48 x float> @llvm.matrix.column.major.load.v48f32.i64(float* nocapture,
 declare void @llvm.matrix.column.major.store.v48f32.i64(<48 x float>, float* nocapture writeonly, i64, i1 immarg, i32 immarg, i32 immarg)
 declare <7 x i1> @llvm.get.active.lane.mask.v7i1.i64(i64, i64)
 declare <7 x float> @llvm.masked.load.v7f32.p0v7f32(<7 x float>*, i32 immarg, <7 x i1>, <7 x float>)
-declare void @llvm.masked.store.v7f32.p0v7f32(<7 x float>, <7 x float>*, i32 immarg, <7 x i1>) 
+declare void @llvm.masked.store.v7f32.p0v7f32(<7 x float>, <7 x float>*, i32 immarg, <7 x i1>)
 declare <7 x float> @llvm.masked.gather.v7f32.v7p0f32(<7 x float*>, i32 immarg, <7 x i1>, <7 x float>)
 declare void @llvm.masked.scatter.v7f32.v7p0f32(<7 x float>, <7 x float*>, i32 immarg, <7 x i1>)
 declare <7 x float> @llvm.masked.expandload.v7f32(float*, <7 x i1>, <7 x float>)
@@ -748,3 +757,5 @@ declare <8 x i64> @llvm.vp.fptoui.v8i64.v8f64(<8 x double>, <8 x i1>, i32)
 declare <8 x i64> @llvm.vp.fptosi.v8i64.v8f64(<8 x double>, <8 x i1>, i32)
 declare <8 x i64> @llvm.vp.ptrtoint.v8i64.v8p0i32(<8 x i32*>, <8 x i1>, i32)
 declare <8 x i32*> @llvm.vp.inttoptr.v8p0i32.v8i64(<8 x i64>, <8 x i1>, i32)
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)