[ORC] Add a bootstrap symbols JITDylib to ExecutionSession. (#188172)

The ExecutionSession constructor now creates a "<bootstrap>" JITDylib
and populates it with the bootstrap symbols from the
ExecutorProcessControl object. This allows bootstrap symbols to be
looked up via ExecutionSession::lookup, providing greater consistency
with other JIT symbol lookups.

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -1408,6 +1408,13 @@ public:
   /// Add a symbol name to the SymbolStringPool and return a pointer to it.
   SymbolStringPtr intern(StringRef SymName) { return EPC->intern(SymName); }
 
+  /// Returns a reference to the bootstrap JITDylib.
+  ///
+  /// This is a bare JITDylib that is created for each ExecutionSession and
+  /// populated with the bootstrap symbol definitions provided by the
+  /// ExecutorProcessControl object.
+  JITDylib &getBootstrapJITDylib() { return BootstrapJD; }
+
   /// Set a WaitingOnGraph::Recorder to capture WaitingOnGraph operations.
   ///
   /// This method can be called at most once. If called, it should be called
@@ -1835,6 +1842,7 @@ private:
   std::vector<ResourceManager *> ResourceManagers;
 
   std::vector<JITDylibSP> JDs;
+  JITDylib &BootstrapJD;
   WaitingOnGraph G;
   WaitingOnGraph::OpRecorder *GOpRecorder = nullptr;
 

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -10,6 +10,7 @@
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Config/llvm-config.h"
+#include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/OrcError.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -1576,9 +1577,16 @@ void LookupTask::printDescription(raw_ostream &OS) { OS << "Lookup task"; }
 void LookupTask::run() { LS.continueLookup(Error::success()); }
 
 ExecutionSession::ExecutionSession(std::unique_ptr<ExecutorProcessControl> EPC)
-    : EPC(std::move(EPC)) {
+    : EPC(std::move(EPC)), BootstrapJD(createBareJITDylib("<bootstrap>")) {
   // Associated EPC and this.
   this->EPC->ES = this;
+  SymbolMap BootstrapSymbols;
+  for (auto &[Name, Ptr] : this->EPC->getBootstrapSymbolsMap())
+    BootstrapSymbols[intern(Name)] =
+        ExecutorSymbolDef(Ptr, JITSymbolFlags::Exported);
+  // Can't fail: BootstrapJD is a new, empty JD and the BootstrapSymbols
+  // variable is a map, so can't contain duplicates.
+  cantFail(BootstrapJD.define(absoluteSymbols(std::move(BootstrapSymbols))));
 }
 
 ExecutionSession::~ExecutionSession() {

llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -1904,4 +1904,67 @@ TEST(CoreAPIsExtraTest, SessionTeardownByFailedToMaterialize) {
   logAllUnhandledErrors(std::move(Err), nulls(), "");
 }
 
+TEST(BootstrapJITDylibTest, BootstrapSymbolsAvailableViaLookup) {
+  // Check that bootstrap symbols provided by the ExecutorProcessControl object
+  // are available via lookup on the bootstrap JITDylib.
+
+  // Create an EPC with some bootstrap symbols pre-populated.
+  auto SSP = std::make_shared<SymbolStringPool>();
+  auto EPC =
+      std::make_unique<UnsupportedExecutorProcessControl>(std::move(SSP));
+
+  // Access the protected BootstrapSymbols map via the public
+  // getBootstrapSymbolsMap accessor... we can't. Instead, use a subclass.
+  class EPCWithBootstrapSymbols : public UnsupportedExecutorProcessControl {
+  public:
+    EPCWithBootstrapSymbols(std::shared_ptr<SymbolStringPool> SSP,
+                            StringMap<ExecutorAddr> BS)
+        : UnsupportedExecutorProcessControl(std::move(SSP)) {
+      this->BootstrapSymbols = std::move(BS);
+    }
+  };
+
+  constexpr ExecutorAddr Addr1(1);
+  constexpr ExecutorAddr Addr2(2);
+
+  StringMap<ExecutorAddr> BootstrapSyms;
+  BootstrapSyms["__orc_rt_run_program"] = Addr1;
+  BootstrapSyms["__orc_rt_log"] = Addr2;
+
+  auto SSP2 = std::make_shared<SymbolStringPool>();
+  auto EPC2 =
+      std::make_unique<EPCWithBootstrapSymbols>(SSP2, std::move(BootstrapSyms));
+  ExecutionSession ES(std::move(EPC2));
+
+  auto &BootstrapJD = ES.getBootstrapJITDylib();
+  EXPECT_EQ(BootstrapJD.getName(), "<bootstrap>");
+
+  // Look up the bootstrap symbols.
+  auto Result =
+      cantFail(ES.lookup(makeJITDylibSearchOrder(&BootstrapJD),
+                         SymbolLookupSet({ES.intern("__orc_rt_run_program"),
+                                          ES.intern("__orc_rt_log")})));
+
+  EXPECT_EQ(Result.size(), 2U);
+  EXPECT_EQ(Result[ES.intern("__orc_rt_run_program")].getAddress(), Addr1);
+  EXPECT_EQ(Result[ES.intern("__orc_rt_log")].getAddress(), Addr2);
+
+  cantFail(ES.endSession());
+}
+
+TEST(BootstrapJITDylibTest, EmptyBootstrapSymbols) {
+  // Check that the bootstrap JITDylib is created even when there are no
+  // bootstrap symbols.
+  auto SSP = std::make_shared<SymbolStringPool>();
+  auto EPC =
+      std::make_unique<UnsupportedExecutorProcessControl>(std::move(SSP));
+
+  ExecutionSession ES(std::move(EPC));
+
+  auto &BootstrapJD = ES.getBootstrapJITDylib();
+  EXPECT_EQ(BootstrapJD.getName(), "<bootstrap>");
+
+  cantFail(ES.endSession());
+}
+
 } // namespace