[ORC] Add a getRequestedSymbols method to MaterializationResponsibility.

This method returns the set of symbols in the target VSO that have queries
waiting on them. This can be used to make decisions about which symbols to
delegate to another MaterializationUnit (typically this will involve
delegating all symbols that have *not* been requested to another
MaterializationUnit so that materialization of those symbols can be
deferred until they are requested).

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@333684 91177308-0d34-0410-b5e6-96231b3b80d8

diff --git a/include/llvm/ExecutionEngine/Orc/Core.h b/include/llvm/ExecutionEngine/Orc/Core.h
index 2b71484..839e88c 100644 (file)
--- a/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/include/llvm/ExecutionEngine/Orc/Core.h
@@ -116,6 +116,12 @@ public:
   ///        into.
   const VSO &getTargetVSO() const { return V; }
 
+  /// Returns the names of any symbols covered by this
+  /// MaterializationResponsibility object that have queries pending. This
+  /// information can be used to return responsibility for unrequested symbols
+  /// back to the VSO via the delegate method.
+  SymbolNameSet getRequestedSymbols();
+
   /// Resolves the given symbols. Individual calls to this method may
   ///        resolve a subset of the symbols, but all symbols must have been
   ///        resolved prior to calling finalize.
@@ -575,6 +581,8 @@ private:
 
   void replace(std::unique_ptr<MaterializationUnit> MU);
 
+  SymbolNameSet getRequestedSymbols(const SymbolFlagsMap &SymbolFlags);
+
   void addDependencies(const SymbolFlagsMap &Dependents,
                        const SymbolDependenceMap &Dependencies);
 

diff --git a/lib/ExecutionEngine/Orc/Core.cpp b/lib/ExecutionEngine/Orc/Core.cpp
index 2fb2e5f..ffa0b17 100644 (file)
--- a/lib/ExecutionEngine/Orc/Core.cpp
+++ b/lib/ExecutionEngine/Orc/Core.cpp
@@ -237,6 +237,10 @@ MaterializationResponsibility::~MaterializationResponsibility() {
          "All symbols should have been explicitly materialized or failed");
 }
 
+SymbolNameSet MaterializationResponsibility::getRequestedSymbols() {
+  return V.getRequestedSymbols(SymbolFlags);
+}
+
 void MaterializationResponsibility::resolve(const SymbolMap &Symbols) {
   for (auto &KV : Symbols) {
     auto I = SymbolFlags.find(KV.first);
@@ -368,16 +372,18 @@ void VSO::replace(std::unique_ptr<MaterializationUnit> MU) {
                  "Can not replace symbol that is not materializing");
           assert(UnmaterializedInfos.count(KV.first) == 0 &&
                  "Symbol being replaced should have no UnmaterializedInfo");
-          assert(MaterializingInfos.count(KV.first) &&
-                 "Symbol being replaced should have a MaterializingInfo");
         }
 #endif // NDEBUG
 
         // If any symbol has pending queries against it then we need to
         // materialize MU immediately.
-        for (auto &KV : MU->getSymbols())
-          if (!MaterializingInfos[KV.first].PendingQueries.empty())
-            return std::move(MU);
+        for (auto &KV : MU->getSymbols()) {
+          auto MII = MaterializingInfos.find(KV.first);
+          if (MII != MaterializingInfos.end()) {
+            if (!MII->second.PendingQueries.empty())
+              return std::move(MU);
+          }
+        }
 
         // Otherwise, make MU responsible for all the symbols.
         auto UMI = std::make_shared<UnmaterializedInfo>(std::move(MU));
@@ -388,8 +394,10 @@ void VSO::replace(std::unique_ptr<MaterializationUnit> MU) {
                  "Materializing flags should be managed internally.");
 
           auto SymI = Symbols.find(KV.first);
-          SymI->second.getFlags() = KV.second;
-          SymI->second.getFlags() |= JITSymbolFlags::Lazy;
+          JITSymbolFlags ReplaceFlags = KV.second;
+          ReplaceFlags |= JITSymbolFlags::Lazy;
+          SymI->second = JITEvaluatedSymbol(SymI->second.getAddress(),
+                                            std::move(ReplaceFlags));
           UnmaterializedInfos[KV.first] = UMI;
         }
 
@@ -400,6 +408,27 @@ void VSO::replace(std::unique_ptr<MaterializationUnit> MU) {
     ES.dispatchMaterialization(*this, std::move(MustRunMU));
 }
 
+SymbolNameSet VSO::getRequestedSymbols(const SymbolFlagsMap &SymbolFlags) {
+  return ES.runSessionLocked([&]() {
+    SymbolNameSet RequestedSymbols;
+
+    for (auto &KV : SymbolFlags) {
+      assert(Symbols.count(KV.first) && "VSO does not cover this symbol?");
+      assert(Symbols[KV.first].getFlags().isMaterializing() &&
+             "getRequestedSymbols can only be called for materializing "
+             "symbols");
+      auto I = MaterializingInfos.find(KV.first);
+      if (I == MaterializingInfos.end())
+        continue;
+
+      if (!I->second.PendingQueries.empty())
+        RequestedSymbols.insert(KV.first);
+    }
+
+    return RequestedSymbols;
+  });
+}
+
 void VSO::addDependencies(const SymbolFlagsMap &Dependants,
                           const SymbolDependenceMap &Dependencies) {
   ES.runSessionLocked([&, this]() {
@@ -870,6 +899,7 @@ VSO &ExecutionSession::createVSO(std::string Name) {
 }
 
 Expected<SymbolMap> lookup(const VSO::VSOList &VSOs, SymbolNameSet Names) {
+
 #if LLVM_ENABLE_THREADS
   // In the threaded case we use promises to return the results.
   std::promise<SymbolMap> PromisedResult;

diff --git a/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 1e2a7ca..bc7a362 100644 (file)
--- a/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -678,4 +678,60 @@ TEST(CoreAPIsTest, TestLookupWithThreadedMaterialization) {
 #endif
 }
 
+TEST(CoreAPIsTest, TestGetRequestedSymbolsAndDelegate) {
+  ExecutionSession ES;
+  auto Foo = ES.getSymbolStringPool().intern("foo");
+  auto Bar = ES.getSymbolStringPool().intern("bar");
+
+  JITEvaluatedSymbol FooSym(0xdeadbeef, JITSymbolFlags::Exported);
+  JITEvaluatedSymbol BarSym(0xcafef00d, JITSymbolFlags::Exported);
+
+  SymbolNameSet Names({Foo, Bar});
+
+  bool FooMaterialized = false;
+  bool BarMaterialized = false;
+
+  auto MU = llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Foo, FooSym.getFlags()}, {Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) {
+        auto Requested = R.getRequestedSymbols();
+        EXPECT_EQ(Requested.size(), 1U) << "Expected one symbol requested";
+        EXPECT_EQ(*Requested.begin(), Foo) << "Expected \"Foo\" requested";
+
+        auto NewMU = llvm::make_unique<SimpleMaterializationUnit>(
+            SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
+            [&](MaterializationResponsibility R2) {
+              R2.resolve(SymbolMap({{Bar, BarSym}}));
+              R2.finalize();
+              BarMaterialized = true;
+            });
+
+        R.delegate(std::move(NewMU));
+
+        R.resolve(SymbolMap({{Foo, FooSym}}));
+        R.finalize();
+
+        FooMaterialized = true;
+      });
+
+  auto &V = ES.createVSO("V");
+
+  cantFail(V.define(MU));
+
+  EXPECT_FALSE(FooMaterialized) << "Foo should not be materialized yet";
+  EXPECT_FALSE(BarMaterialized) << "Bar should not be materialized yet";
+
+  auto FooSymResult = cantFail(lookup({&V}, Foo));
+  EXPECT_EQ(FooSymResult.getAddress(), FooSym.getAddress())
+      << "Address mismatch for Foo";
+
+  EXPECT_TRUE(FooMaterialized) << "Foo should be materialized now";
+  EXPECT_FALSE(BarMaterialized) << "Bar still should not be materialized";
+
+  auto BarSymResult = cantFail(lookup({&V}, Bar));
+  EXPECT_EQ(BarSymResult.getAddress(), BarSym.getAddress())
+      << "Address mismatch for Bar";
+  EXPECT_TRUE(BarMaterialized) << "Bar should be materialized now";
+}
+
 } // namespace