[PATCH] Enable using llvm jitlink as an alternative llvm jit linker of old Rtdyld. - Mailing list pgsql-hackers

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. This also adds minor llvm 15 compat fix that is needed
---
 config/llvm.m4                        |  1 +
 src/backend/jit/llvm/llvmjit.c        | 67 +++++++++++++++++++++++++--
 src/backend/jit/llvm/llvmjit_wrap.cpp | 35 ++++++++++++++
 src/include/jit/llvmjit.h             |  9 ++++
 4 files changed, 108 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 6c72d43beb..d8b840da8c 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;
 
@@ -820,16 +861,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;
@@ -1112,7 +1158,11 @@ llvm_resolve_symbols(LLVMOrcDefinitionGeneratorRef GeneratorObj, void *Ctx,
                      LLVMOrcJITDylibRef JD, LLVMOrcJITDylibLookupFlags JDLookupFlags,
                      LLVMOrcCLookupSet LookupSet, size_t LookupSetSize)
 {
+#if LLVM_VERSION_MAJOR > 14
+    LLVMOrcCSymbolMapPairs symbols = palloc0(sizeof(LLVMOrcCSymbolMapPair) * LookupSetSize);
+#else
     LLVMOrcCSymbolMapPairs symbols = palloc0(sizeof(LLVMJITCSymbolMapPair) * LookupSetSize);
+#endif
     LLVMErrorRef error;
     LLVMOrcMaterializationUnitRef mu;
 
@@ -1160,6 +1210,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);
 
@@ -1179,6 +1233,7 @@ llvm_create_object_layer(void *Ctx, LLVMOrcExecutionSessionRef ES, const char *T
 
         LLVMOrcRTDyldObjectLinkingLayerRegisterJITEventListener(objlayer, l);
     }
+#endif
 #endif
 
     return objlayer;
@@ -1230,7 +1285,11 @@ llvm_create_jit_instance(LLVMTargetMachineRef tm)
      * Symbol resolution support for "special" functions, e.g. a call into an
      * SQL callable function.
      */
+#if LLVM_VERSION_MAJOR > 14
+    ref_gen = LLVMOrcCreateCustomCAPIDefinitionGenerator(llvm_resolve_symbols, NULL, NULL);
+#else
     ref_gen = LLVMOrcCreateCustomCAPIDefinitionGenerator(llvm_resolve_symbols, NULL);
+#endif
     LLVMOrcJITDylibAddGenerator(LLVMOrcLLJITGetMainJITDylib(lljit), ref_gen);
 
     return lljit;
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.2