From fd5b3ac927fb40155dc0ebe5bc5ec83d4ca5ca4f Mon Sep 17 00:00:00 2001
From: Alex Fan <alex.fan.q@gmail.com>
Date: Mon, 29 Aug 2022 15:24:16 +0800
Subject: [PATCH] Enable using llvm jitlink as an alternative llvm jit linker
 of old Rtdyld.

This brings the bonus of support jitting on riscv64 (included in this patch)
and other platforms Rtdyld doesn't support, e.g. windows COFF.

Currently, llvm doesn't expose jitlink (ObjectLinkingLayer) via C API, so
a wrapper is added.
---
 config/llvm.m4                        |  1 +
 src/backend/jit/llvm/llvmjit.c        | 59 +++++++++++++++++++++++++--
 src/backend/jit/llvm/llvmjit_wrap.cpp | 35 ++++++++++++++++
 src/include/jit/llvmjit.h             |  9 ++++
 4 files changed, 100 insertions(+), 4 deletions(-)

diff --git a/config/llvm.m4 b/config/llvm.m4
index 3a75cd8b4d..a31b8b304a 100644
--- a/config/llvm.m4
+++ b/config/llvm.m4
@@ -75,6 +75,7 @@ AC_DEFUN([PGAC_LLVM_SUPPORT],
       engine) pgac_components="$pgac_components $pgac_component";;
       debuginfodwarf) pgac_components="$pgac_components $pgac_component";;
       orcjit) pgac_components="$pgac_components $pgac_component";;
+      jitlink) pgac_components="$pgac_components $pgac_component";;
       passes) pgac_components="$pgac_components $pgac_component";;
       native) pgac_components="$pgac_components $pgac_component";;
       perfjitevents) pgac_components="$pgac_components $pgac_component";;
diff --git a/src/backend/jit/llvm/llvmjit.c b/src/backend/jit/llvm/llvmjit.c
index 9aca7fc7a4..a444b89a65 100644
--- a/src/backend/jit/llvm/llvmjit.c
+++ b/src/backend/jit/llvm/llvmjit.c
@@ -229,6 +229,11 @@ llvm_release_context(JitContext *context)
 LLVMModuleRef
 llvm_mutable_module(LLVMJitContext *context)
 {
+#ifdef __riscv
+	const char* abiname;
+	const char* target_abi = "target-abi";
+	LLVMMetadataRef abi_metadata;
+#endif
 	llvm_assert_in_fatal_section();
 
 	/*
@@ -241,6 +246,40 @@ llvm_mutable_module(LLVMJitContext *context)
 		context->module = LLVMModuleCreateWithName("pg");
 		LLVMSetTarget(context->module, llvm_triple);
 		LLVMSetDataLayout(context->module, llvm_layout);
+#ifdef __riscv
+#if __riscv_xlen == 64
+#ifdef __riscv_float_abi_double
+		abiname = "lp64d";
+#elif defined(__riscv_float_abi_single)
+		abiname = "lp64f";
+#else
+		abiname = "lp64";
+#endif
+#elif __riscv_xlen == 32
+#ifdef __riscv_float_abi_double
+		abiname = "ilp32d";
+#elif defined(__riscv_float_abi_single)
+		abiname = "ilp32f";
+#else
+		abiname = "ilp32";
+#endif
+#else
+		elog(ERROR, "unsupported riscv xlen %d", __riscv_xlen);
+#endif
+		/*
+		 * set this manually to avoid llvm defaulting to soft float and
+		 * resulting in linker error: `can't link double-float modules
+		 * with soft-float modules`
+		 * we could set this for TargetMachine via MCOptions, but there
+		 * is no C API for it
+		 * ref: https://github.com/llvm/llvm-project/blob/afa520ab34803c82587ea6759bfd352579f741b4/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp#L90
+		 */
+		abi_metadata = LLVMMDStringInContext2(
+			LLVMGetModuleContext(context->module),
+			abiname, strlen(abiname));
+		LLVMAddModuleFlag(context->module, LLVMModuleFlagBehaviorOverride,
+			target_abi, strlen(target_abi), abi_metadata);
+#endif
 	}
 
 	return context->module;
@@ -786,6 +825,8 @@ llvm_session_initialize(void)
 	char	   *error = NULL;
 	char	   *cpu = NULL;
 	char	   *features = NULL;
+	LLVMRelocMode reloc=LLVMRelocDefault;
+	LLVMCodeModel codemodel=LLVMCodeModelJITDefault;
 	LLVMTargetMachineRef opt0_tm;
 	LLVMTargetMachineRef opt3_tm;
 
@@ -830,16 +871,21 @@ llvm_session_initialize(void)
 	elog(DEBUG2, "LLVMJIT detected CPU \"%s\", with features \"%s\"",
 		 cpu, features);
 
+#ifdef __riscv
+	reloc=LLVMRelocPIC;
+	codemodel=LLVMCodeModelMedium;
+#endif
+
 	opt0_tm =
 		LLVMCreateTargetMachine(llvm_targetref, llvm_triple, cpu, features,
 								LLVMCodeGenLevelNone,
-								LLVMRelocDefault,
-								LLVMCodeModelJITDefault);
+								reloc,
+								codemodel);
 	opt3_tm =
 		LLVMCreateTargetMachine(llvm_targetref, llvm_triple, cpu, features,
 								LLVMCodeGenLevelAggressive,
-								LLVMRelocDefault,
-								LLVMCodeModelJITDefault);
+								reloc,
+								codemodel);
 
 	LLVMDisposeMessage(cpu);
 	cpu = NULL;
@@ -1174,6 +1220,10 @@ llvm_log_jit_error(void *ctx, LLVMErrorRef error)
 static LLVMOrcObjectLayerRef
 llvm_create_object_layer(void *Ctx, LLVMOrcExecutionSessionRef ES, const char *Triple)
 {
+#if defined(USE_JITLINK)
+	LLVMOrcObjectLayerRef objlayer =
+	LLVMOrcCreateJitlinkObjectLinkingLayer(ES);
+#else
 	LLVMOrcObjectLayerRef objlayer =
 	LLVMOrcCreateRTDyldObjectLinkingLayerWithSectionMemoryManager(ES);
 
@@ -1193,6 +1243,7 @@ llvm_create_object_layer(void *Ctx, LLVMOrcExecutionSessionRef ES, const char *T
 
 		LLVMOrcRTDyldObjectLinkingLayerRegisterJITEventListener(objlayer, l);
 	}
+#endif
 #endif
 
 	return objlayer;
diff --git a/src/backend/jit/llvm/llvmjit_wrap.cpp b/src/backend/jit/llvm/llvmjit_wrap.cpp
index 8f11cc02b2..29f21f1715 100644
--- a/src/backend/jit/llvm/llvmjit_wrap.cpp
+++ b/src/backend/jit/llvm/llvmjit_wrap.cpp
@@ -27,6 +27,10 @@ extern "C"
 #include <llvm/Support/Host.h>
 
 #include "jit/llvmjit.h"
+#ifdef USE_JITLINK
+#include "llvm/ExecutionEngine/JITLink/EHFrameSupport.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#endif
 
 
 /*
@@ -48,6 +52,19 @@ char *LLVMGetHostCPUFeatures(void) {
 		for (auto &F : HostFeatures)
 			Features.AddFeature(F.first(), F.second);
 
+#if defined(__riscv)
+	/* getHostCPUName returns "generic-rv[32|64]", which lacks all features */
+	Features.AddFeature("m", true);
+	Features.AddFeature("a", true);
+	Features.AddFeature("c", true);
+# if defined(__riscv_float_abi_single)
+	Features.AddFeature("f", true);
+# endif
+# if defined(__riscv_float_abi_double)
+	Features.AddFeature("d", true);
+# endif
+#endif
+
 	return strdup(Features.getString().c_str());
 }
 #endif
@@ -76,3 +93,21 @@ LLVMGetAttributeCountAtIndexPG(LLVMValueRef F, uint32 Idx)
 	 */
 	return LLVMGetAttributeCountAtIndex(F, Idx);
 }
+
+#ifdef USE_JITLINK
+/*
+ * There is no public C API to create ObjectLinkingLayer for JITLINK, create our own
+ */
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(llvm::orc::ExecutionSession, LLVMOrcExecutionSessionRef)
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(llvm::orc::ObjectLayer, LLVMOrcObjectLayerRef)
+
+LLVMOrcObjectLayerRef
+LLVMOrcCreateJitlinkObjectLinkingLayer(LLVMOrcExecutionSessionRef ES)
+{
+	assert(ES && "ES must not be null");
+	auto ObjLinkingLayer = new llvm::orc::ObjectLinkingLayer(*unwrap(ES));
+	ObjLinkingLayer->addPlugin(std::make_unique<llvm::orc::EHFrameRegistrationPlugin>(
+		*unwrap(ES), std::make_unique<llvm::jitlink::InProcessEHFrameRegistrar>()));
+	return wrap(ObjLinkingLayer);
+}
+#endif
diff --git a/src/include/jit/llvmjit.h b/src/include/jit/llvmjit.h
index 4541f9a2c4..85a0cfe5e0 100644
--- a/src/include/jit/llvmjit.h
+++ b/src/include/jit/llvmjit.h
@@ -19,6 +19,11 @@
 
 #include <llvm-c/Types.h>
 
+#if defined(__riscv) && LLVM_VERSION_MAJOR >= 15
+#include <llvm-c/Orc.h>
+#define USE_JITLINK
+/* else use legacy RTDyld */
+#endif
 
 /*
  * File needs to be includable by both C and C++ code, and include other
@@ -134,6 +139,10 @@ extern char *LLVMGetHostCPUFeatures(void);
 
 extern unsigned LLVMGetAttributeCountAtIndexPG(LLVMValueRef F, uint32 Idx);
 
+#ifdef USE_JITLINK
+extern LLVMOrcObjectLayerRef LLVMOrcCreateJitlinkObjectLinkingLayer(LLVMOrcExecutionSessionRef ES);
+#endif
+
 #ifdef __cplusplus
 } /* extern "C" */
 #endif
-- 
2.37.3

