add uncheckedGetIntrinsicSignature and mlirLLVMFunctionTypeGetReturnType

Author: makslevental

mlir/include/mlir-c/Dialect/LLVM.h

@@ -52,6 +52,9 @@ MLIR_CAPI_EXPORTED intptr_t mlirLLVMFunctionTypeGetNumInputs(MlirType type);
 MLIR_CAPI_EXPORTED MlirType mlirLLVMFunctionTypeGetInput(MlirType type,
                                                          intptr_t pos);
 
+/// Returns the return type of the function type.
+MLIR_CAPI_EXPORTED MlirType mlirLLVMFunctionTypeGetReturnType(MlirType type);
+
 /// Returns `true` if the type is an LLVM dialect struct type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsALLVMStructType(MlirType type);
 

mlir/include/mlir-c/Target/LLVMIR.h

@@ -74,6 +74,14 @@ mlirTypeToLLVMIRTranslatorDestroy(MlirTypeToLLVMIRTranslator translator);
 MLIR_CAPI_EXPORTED LLVMTypeRef mlirTypeToLLVMIRTranslatorTranslateType(
     MlirTypeToLLVMIRTranslator translator, MlirType mlirType);
 
+/// Attempts to look up the LLVM intrinsic matching `id` with parameter types
+/// `mlirParamTypes`. Note, since no return type is required, no check is
+/// performed to verify the found intrinsic.
+MLIR_CAPI_EXPORTED LLVMTypeRef
+mlirTypeToLLVMIRTranslatorUncheckedGetIntrinsicSignature(
+    MlirTypeToLLVMIRTranslator translator, unsigned id,
+    MlirType *mlirParamTypes, unsigned numParams);
+
 #ifdef __cplusplus
 }
 #endif

mlir/include/mlir/Target/LLVMIR/TypeToLLVM.h

@@ -14,9 +14,12 @@
 #ifndef MLIR_TARGET_LLVMIR_TYPETOLLVM_H
 #define MLIR_TARGET_LLVMIR_TYPETOLLVM_H
 
+#include "llvm/ADT/ArrayRef.h"
+
 #include <memory>
 
 namespace llvm {
+class FunctionType;
 class DataLayout;
 class LLVMContext;
 class Type;
@@ -50,6 +53,12 @@ class TypeToLLVMIRTranslator {
   /// Translates the given MLIR LLVM dialect type to LLVM IR.
   llvm::Type *translateType(Type type);
 
+  /// Attempts to look up the LLVM intrinsic matching `id` with parameter types
+  /// `types`. Note, since no return type is required, no check is performed to
+  /// verify the found intrinsic.
+  llvm::FunctionType *
+  uncheckedGetIntrinsicSignature(unsigned id, llvm::ArrayRef<Type> types);
+
 private:
   /// Private implementation.
   std::unique_ptr<detail::TypeToLLVMIRTranslatorImpl> impl;

mlir/lib/CAPI/Dialect/LLVM.cpp

@@ -65,6 +65,10 @@ MlirType mlirLLVMFunctionTypeGetInput(MlirType type, intptr_t pos) {
                   .getParamType(static_cast<unsigned>(pos)));
 }
 
+MlirType mlirLLVMFunctionTypeGetReturnType(MlirType type) {
+  return wrap(llvm::cast<LLVM::LLVMFunctionType>(unwrap(type)).getReturnType());
+}
+
 bool mlirTypeIsALLVMStructType(MlirType type) {
   return isa<LLVM::LLVMStructType>(unwrap(type));
 }

mlir/lib/CAPI/Target/LLVMIR.cpp

@@ -9,6 +9,7 @@
 
 #include "mlir-c/Target/LLVMIR.h"
 
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
@@ -77,3 +78,15 @@ mlirTypeToLLVMIRTranslatorTranslateType(MlirTypeToLLVMIRTranslator translator,
   llvm::Type *type = translator_->translateType(unwrap(mlirType));
   return llvm::wrap(type);
 }
+
+LLVMTypeRef mlirTypeToLLVMIRTranslatorUncheckedGetIntrinsicSignature(
+    MlirTypeToLLVMIRTranslator translator, unsigned id,
+    MlirType *mlirParamTypes, unsigned numParams) {
+  SmallVector<mlir::Type, 4> paramTypes;
+  for (unsigned i = 0; i < numParams; i++)
+    paramTypes.emplace_back(unwrap(mlirParamTypes[i]));
+  LLVM::TypeToLLVMIRTranslator *translator_ = unwrap(translator);
+  llvm::Type *functionType =
+      translator_->uncheckedGetIntrinsicSignature(id, paramTypes);
+  return llvm::wrap(functionType);
+}

mlir/lib/Target/LLVMIR/TypeToLLVM.cpp

@@ -14,6 +14,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Type.h"
 
 using namespace mlir;
@@ -84,6 +85,15 @@ class TypeToLLVMIRTranslatorImpl {
     return translated;
   }
 
+  llvm::FunctionType *
+  uncheckedGetIntrinsicSignature(unsigned id, llvm::ArrayRef<Type> types) {
+    SmallVector<llvm::Type *, 8> paramTypes;
+    if (llvm::Intrinsic::isOverloaded(id)) {
+      translateTypes(types, paramTypes);
+    }
+    return llvm::Intrinsic::getType(context, id, paramTypes);
+  }
+
 private:
   /// Translates the given array type.
   llvm::Type *translate(LLVM::LLVMArrayType type) {
@@ -193,6 +203,12 @@ llvm::Type *LLVM::TypeToLLVMIRTranslator::translateType(Type type) {
   return impl->translateType(type);
 }
 
+llvm::FunctionType *
+LLVM::TypeToLLVMIRTranslator::uncheckedGetIntrinsicSignature(
+    unsigned id, llvm::ArrayRef<Type> types) {
+  return impl->uncheckedGetIntrinsicSignature(id, types);
+}
+
 unsigned LLVM::TypeToLLVMIRTranslator::getPreferredAlignment(
     Type type, const llvm::DataLayout &layout) {
   return layout.getPrefTypeAlign(translateType(type)).value();

mlir/test/CAPI/translation.c

@@ -58,9 +58,9 @@ static void testToLLVMIR(MlirContext ctx) {
   LLVMContextDispose(llvmCtx);
 }
 
-// CHECK-LABEL: testTypeFromLLVMIRTranslator
-static void testTypeFromLLVMIRTranslator(MlirContext ctx) {
-  fprintf(stderr, "testTypeFromLLVMIRTranslator\n");
+// CHECK-LABEL: testTypeToFromLLVMIRTranslator
+static void testTypeToFromLLVMIRTranslator(MlirContext ctx) {
+  fprintf(stderr, "testTypeToFromLLVMIRTranslator\n");
   LLVMContextRef llvmCtx = LLVMContextCreate();
 
   LLVMTypeRef llvmTy = LLVMInt32TypeInContext(llvmCtx);
@@ -77,6 +77,37 @@ static void testTypeFromLLVMIRTranslator(MlirContext ctx) {
       mlirTypeToLLVMIRTranslatorTranslateType(toLLVMTranslator, mlirTy);
   // CHECK: i32
   LLVMDumpType(llvmTy2);
+  fprintf(stderr, "\n");
+
+  // check "not overloaded" path
+  MlirType i32 = mlirIntegerTypeGet(ctx, 32);
+  MlirType nonOverloadeTys[] = {i32};
+  unsigned returnAddressIID = LLVMLookupIntrinsicID("llvm.returnaddress", 18);
+  LLVMTypeRef returnAddressIntrinsicTy =
+      mlirTypeToLLVMIRTranslatorUncheckedGetIntrinsicSignature(
+          toLLVMTranslator, returnAddressIID, nonOverloadeTys, 1);
+  // CHECK: ptr (i32)
+  LLVMDumpType(returnAddressIntrinsicTy);
+  fprintf(stderr, "\n");
+
+  // check "overloaded" path
+  MlirType overloadeTys[] = {i32, i32};
+  unsigned sMaxIID = LLVMLookupIntrinsicID("llvm.smax", 9);
+  LLVMTypeRef llvmSMaxIntrinsicTy =
+      mlirTypeToLLVMIRTranslatorUncheckedGetIntrinsicSignature(
+          toLLVMTranslator, sMaxIID, overloadeTys, 2);
+  // CHECK: i32 (i32, i32)
+  LLVMDumpType(llvmSMaxIntrinsicTy);
+  fprintf(stderr, "\n");
+
+  MlirType mlirSMaxIntrinsicTy = mlirTypeFromLLVMIRTranslatorTranslateType(
+      fromLLVMTranslator, llvmSMaxIntrinsicTy);
+  // CHECK: !llvm.func<i32 (i32, i32)>
+  mlirTypeDump(mlirSMaxIntrinsicTy);
+
+  MlirType retType = mlirLLVMFunctionTypeGetReturnType(mlirSMaxIntrinsicTy);
+  // CHECK: i32
+  mlirTypeDump(retType);
 
   mlirTypeFromLLVMIRTranslatorDestroy(fromLLVMTranslator);
   mlirTypeToLLVMIRTranslatorDestroy(toLLVMTranslator);
@@ -88,7 +119,7 @@ int main(void) {
   mlirDialectHandleRegisterDialect(mlirGetDialectHandle__llvm__(), ctx);
   mlirContextGetOrLoadDialect(ctx, mlirStringRefCreateFromCString("llvm"));
   testToLLVMIR(ctx);
-  testTypeFromLLVMIRTranslator(ctx);
+  testTypeToFromLLVMIRTranslator(ctx);
   mlirContextDestroy(ctx);
   return 0;
 }