From fd37a40f2db0706f94089ef42f6ca30da810d1bb Mon Sep 17 00:00:00 2001
From: Edd Barrett <vext01@gmail.com>
Date: Wed, 21 Sep 2022 16:49:10 +0100
Subject: [PATCH] Add the `-yk-alloc-llvmbc-section` flag.

Enabling this flag does three things:

 - Adds the `SHF_ALLOC` flag to the `.llvmbc` section so that it gets
   loaded by the loader.

 - Makes the `llvm.embedded.module` symbol inside externally visible by
   giving it external linkage.

 - Prepends the embedded IR with a pointer-sized length field so that we
   can load the IR directly from memory at runtime.

I was having issues finding a place to define the flag where it could be
consistently accessed across the necessary module linkage boundaries.
I've ended up defining  it in `llvm/lib/Support`, which is an "always
linked" component.
---
 llvm/lib/Bitcode/Writer/BitcodeWriter.cpp     | 27 +++++++++++++++++--
 .../CodeGen/TargetLoweringObjectFileImpl.cpp  |  8 ++++++
 llvm/lib/Support/CMakeLists.txt               |  1 +
 llvm/lib/Support/Yk.cpp                       |  8 ++++++
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 llvm/lib/Support/Yk.cpp

diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index d7e012fb6a9ed..8dea754c1478d 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -98,6 +98,8 @@ static cl::opt<bool> WriteRelBFToSummary(
 
 extern FunctionSummary::ForceSummaryHotnessType ForceSummaryEdgesCold;
 
+extern bool YkAllocLLVMBCSection;
+
 namespace {
 
 /// These are manifest constants used by the bitcode writer. They do not need to
@@ -4992,11 +4994,32 @@ void llvm::embedBitcodeInModule(llvm::Module &M, llvm::MemoryBufferRef Buf,
       ModuleData = ArrayRef<uint8_t>((const uint8_t *)Buf.getBufferStart(),
                                      Buf.getBufferSize());
   }
+
+  GlobalValue::LinkageTypes SymLinkage = GlobalValue::PrivateLinkage;
+
+  // For the YK JIT we prepend a header containing the size of the bitcode.
+  // This is required in order to load bitcode from memory.
+  std::vector<uint8_t> YkModuleData;
+  if (YkAllocLLVMBCSection) {
+    // Write length field.
+    size_t ModuleDataSize = ModuleData.size();
+    uint8_t *Bytes = reinterpret_cast<uint8_t *>(&ModuleDataSize);
+    for (size_t I = 0; I < sizeof(ModuleDataSize); I++)
+      YkModuleData.push_back(Bytes[I]);
+
+    // Append bitcode.
+    std::move(ModuleData.begin(), ModuleData.end(),
+              std::back_inserter(YkModuleData));
+    ModuleData = YkModuleData;
+
+    // Ensure the symbol is exported in the resulting binary.
+    SymLinkage = GlobalValue::ExternalLinkage;
+  }
+
   llvm::Constant *ModuleConstant =
       llvm::ConstantDataArray::get(M.getContext(), ModuleData);
   llvm::GlobalVariable *GV = new llvm::GlobalVariable(
-      M, ModuleConstant->getType(), true, llvm::GlobalValue::PrivateLinkage,
-      ModuleConstant);
+      M, ModuleConstant->getType(), true, SymLinkage, ModuleConstant);
   GV->setSection(getSectionNameForBitcode(T));
   // Set alignment to 1 to prevent padding between two contributions from input
   // sections after linking.
diff --git a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
index 2badbe34ae6a7..1c87d2d3d8728 100644
--- a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
@@ -70,6 +70,8 @@
 using namespace llvm;
 using namespace dwarf;
 
+extern bool YkAllocLLVMBCSection;
+
 static void GetObjCImageInfo(Module &M, unsigned &Version, unsigned &Flags,
                              StringRef &Section) {
   SmallVector<Module::ModuleFlagEntry, 8> ModuleFlags;
@@ -786,6 +788,12 @@ static MCSection *selectExplicitSectionGlobal(
       Retain, ForceUnique);
 
   const MCSymbolELF *LinkedToSym = getLinkedToSymbol(GO, TM);
+
+  // The Yk JIT expects to load the IR from its address space. This tells the
+  // loader to load the section.
+  if (YkAllocLLVMBCSection && (SectionName == ".llvmbc"))
+    Flags |= llvm::ELF::SHF_ALLOC;
+
   MCSectionELF *Section = Ctx.getELFSection(
       SectionName, getELFSectionType(SectionName, Kind), Flags, EntrySize,
       Group, IsComdat, UniqueID, LinkedToSym);
diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
index 5044b2639a0f7..2cf5565c05fb5 100644
--- a/llvm/lib/Support/CMakeLists.txt
+++ b/llvm/lib/Support/CMakeLists.txt
@@ -236,6 +236,7 @@ add_llvm_component_library(LLVMSupport
   X86TargetParser.cpp
   YAMLParser.cpp
   YAMLTraits.cpp
+  Yk.cpp
   raw_os_ostream.cpp
   raw_ostream.cpp
   regcomp.c
diff --git a/llvm/lib/Support/Yk.cpp b/llvm/lib/Support/Yk.cpp
new file mode 100644
index 0000000000000..056243b45ce85
--- /dev/null
+++ b/llvm/lib/Support/Yk.cpp
@@ -0,0 +1,8 @@
+#include "llvm/Support/CommandLine.h"
+
+using namespace llvm;
+
+bool YkAllocLLVMBCSection;
+static cl::opt<bool, true> YkAllocLLVMBCSectionParser(
+    "yk-alloc-llvmbc-section", cl::desc("Make the `.llvmbc` section loadable"),
+    cl::NotHidden, cl::location(YkAllocLLVMBCSection));