[llvm] 14b7c10 - [C-API][ORC] Add C API to suspend lookups during definition generation.

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Mon Jun 13 17:50:20 PDT 2022


Author: Lang Hames
Date: 2022-06-13T17:20:07-07:00
New Revision: 14b7c108a2bf46541efc3a5c9cbd589b3afc18e6

URL: https://github.com/llvm/llvm-project/commit/14b7c108a2bf46541efc3a5c9cbd589b3afc18e6
DIFF: https://github.com/llvm/llvm-project/commit/14b7c108a2bf46541efc3a5c9cbd589b3afc18e6.diff

LOG: [C-API][ORC] Add C API to suspend lookups during definition generation.

Slow definition generators may suspend lookups to temporarily release the
session lock, allowing unrelated lookups to proceed.

Using this functionality is discouraged: it is best to make definition
generation fast, rather than suspending the lookup. As a last resort where
this is not possible, suspension may be used.

Added: 
    

Modified: 
    llvm/include/llvm-c/Orc.h
    llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
    llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm-c/Orc.h b/llvm/include/llvm-c/Orc.h
index e843e76a6da0..0dcfb06865aa 100644
--- a/llvm/include/llvm-c/Orc.h
+++ b/llvm/include/llvm-c/Orc.h
@@ -357,6 +357,14 @@ typedef LLVMErrorRef (*LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction)(
     LLVMOrcJITDylibRef JD, LLVMOrcJITDylibLookupFlags JDLookupFlags,
     LLVMOrcCLookupSet LookupSet, size_t LookupSetSize);
 
+/**
+ * Disposer for a custom generator.
+ *
+ * Will be called by ORC when the JITDylib that the generator is attached to
+ * is destroyed.
+ */
+typedef void (*LLVMOrcDisposeCAPIDefinitionGeneratorFunction)(void *Ctx);
+
 /**
  * Predicate function for SymbolStringPoolEntries.
  */
@@ -573,6 +581,11 @@ void LLVMOrcRetainSymbolStringPoolEntry(LLVMOrcSymbolStringPoolEntryRef S);
  */
 void LLVMOrcReleaseSymbolStringPoolEntry(LLVMOrcSymbolStringPoolEntryRef S);
 
+/**
+ * Return the c-string for the given symbol. This string will remain valid until
+ * the entry is freed (once all LLVMOrcSymbolStringPoolEntryRefs have been
+ * released).
+ */
 const char *LLVMOrcSymbolStringPoolEntryStr(LLVMOrcSymbolStringPoolEntryRef S);
 
 /**
@@ -970,9 +983,27 @@ void LLVMOrcJITDylibAddGenerator(LLVMOrcJITDylibRef JD,
 
 /**
  * Create a custom generator.
+ *
+ * The F argument will be used to implement the DefinitionGenerator's
+ * tryToGenerate method (see
+ * LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction).
+ *
+ * Ctx is a context object that will be passed to F. This argument is
+ * permitted to be null.
+ *
+ * Dispose is the disposal function for Ctx. This argument is permitted to be
+ * null (in which case the client is responsible for the lifetime of Ctx).
  */
 LLVMOrcDefinitionGeneratorRef LLVMOrcCreateCustomCAPIDefinitionGenerator(
-    LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F, void *Ctx);
+    LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F, void *Ctx,
+    LLVMOrcDisposeCAPIDefinitionGeneratorFunction Dispose);
+
+/**
+ * Continue a lookup that was suspended in a generator (see
+ * LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction).
+ */
+void LLVMOrcLookupStateContinueLookup(LLVMOrcLookupStateRef S,
+                                      LLVMErrorRef Err);
 
 /**
  * Get a DynamicLibrarySearchGenerator that will reflect process symbols into

diff  --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
index 0d1024ab1455..b7eab6b85ecf 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
@@ -281,9 +281,14 @@ namespace orc {
 class CAPIDefinitionGenerator final : public DefinitionGenerator {
 public:
   CAPIDefinitionGenerator(
-      void *Ctx,
+      LLVMOrcDisposeCAPIDefinitionGeneratorFunction Dispose, void *Ctx,
       LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction TryToGenerate)
-      : Ctx(Ctx), TryToGenerate(TryToGenerate) {}
+      : Dispose(Dispose), Ctx(Ctx), TryToGenerate(TryToGenerate) {}
+
+  ~CAPIDefinitionGenerator() {
+    if (Dispose)
+      Dispose(Ctx);
+  }
 
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
@@ -321,6 +326,7 @@ class CAPIDefinitionGenerator final : public DefinitionGenerator {
   }
 
 private:
+  LLVMOrcDisposeCAPIDefinitionGeneratorFunction Dispose;
   void *Ctx;
   LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction TryToGenerate;
 };
@@ -669,11 +675,19 @@ void LLVMOrcJITDylibAddGenerator(LLVMOrcJITDylibRef JD,
 }
 
 LLVMOrcDefinitionGeneratorRef LLVMOrcCreateCustomCAPIDefinitionGenerator(
-    LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F, void *Ctx) {
-  auto DG = std::make_unique<CAPIDefinitionGenerator>(Ctx, F);
+    LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F, void *Ctx,
+    LLVMOrcDisposeCAPIDefinitionGeneratorFunction Dispose) {
+  auto DG = std::make_unique<CAPIDefinitionGenerator>(Dispose, Ctx, F);
   return wrap(DG.release());
 }
 
+void LLVMOrcLookupStateContinueLookup(LLVMOrcLookupStateRef S,
+                                      LLVMErrorRef Err) {
+  LookupState LS;
+  OrcV2CAPIHelper::resetLookupState(LS, ::unwrap(S));
+  LS.continueLookup(unwrap(Err));
+}
+
 LLVMErrorRef LLVMOrcCreateDynamicLibrarySearchGeneratorForProcess(
     LLVMOrcDefinitionGeneratorRef *Result, char GlobalPrefix,
     LLVMOrcSymbolPredicate Filter, void *FilterCtx) {

diff  --git a/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp b/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
index 2660c8bf72b2..37afaa0bafe6 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/OrcCAPITest.cpp
@@ -107,8 +107,12 @@ class OrcCAPITestBase : public testing::Test {
     MainDylib = LLVMOrcLLJITGetMainJITDylib(Jit);
   }
   void TearDown() override {
-    LLVMOrcDisposeLLJIT(Jit);
-    Jit = nullptr;
+    // Check whether Jit has already been torn down -- we allow clients to do
+    // this manually to check teardown behavior.
+    if (Jit) {
+      LLVMOrcDisposeLLJIT(Jit);
+      Jit = nullptr;
+    }
   }
 
 protected:
@@ -410,7 +414,7 @@ TEST_F(OrcCAPITestBase, ExecutionSessionLookup_Failure) {
 TEST_F(OrcCAPITestBase, DefinitionGenerators) {
   LLVMOrcDefinitionGeneratorRef Gen =
       LLVMOrcCreateCustomCAPIDefinitionGenerator(&definitionGeneratorFn,
-                                                 nullptr);
+                                                 nullptr, nullptr);
   LLVMOrcJITDylibAddGenerator(MainDylib, Gen);
   LLVMOrcJITTargetAddress OutAddr;
   if (LLVMErrorRef E = LLVMOrcLLJITLookup(Jit, &OutAddr, "test"))
@@ -660,3 +664,110 @@ TEST_F(OrcCAPITestBase, MaterializationResponsibility) {
   ASSERT_TRUE(!!Addr);
   ASSERT_EQ(Addr, (LLVMOrcJITTargetAddress)&TargetFn);
 }
+
+struct SuspendedLookupContext {
+  std::function<void()> AsyncWork;
+  LLVMOrcSymbolStringPoolEntryRef NameToGenerate;
+  JITTargetAddress AddrToGenerate;
+
+  bool Disposed = false;
+  bool QueryCompleted = true;
+};
+
+static LLVMErrorRef TryToGenerateWithSuspendedLookup(
+    LLVMOrcDefinitionGeneratorRef GeneratorObj, void *RawCtx,
+    LLVMOrcLookupStateRef *LookupState, LLVMOrcLookupKind Kind,
+    LLVMOrcJITDylibRef JD, LLVMOrcJITDylibLookupFlags JDLookupFlags,
+    LLVMOrcCLookupSet LookupSet, size_t LookupSetSize) {
+
+  auto *Ctx = static_cast<SuspendedLookupContext *>(RawCtx);
+
+  assert(LookupSetSize == 1);
+  assert(LookupSet[0].Name == Ctx->NameToGenerate);
+
+  LLVMJITEvaluatedSymbol Sym = {0x1234, {LLVMJITSymbolGenericFlagsExported, 0}};
+  LLVMOrcRetainSymbolStringPoolEntry(LookupSet[0].Name);
+  LLVMOrcCSymbolMapPair Pair = {LookupSet[0].Name, Sym};
+  LLVMOrcCSymbolMapPair Pairs[] = {Pair};
+  LLVMOrcMaterializationUnitRef MU = LLVMOrcAbsoluteSymbols(Pairs, 1);
+
+  // Capture and reset LookupState to suspend the lookup. We'll continue it in
+  // the SuspendedLookup testcase below.
+  Ctx->AsyncWork = [LS = *LookupState, JD, MU]() {
+    LLVMErrorRef Err = LLVMOrcJITDylibDefine(JD, MU);
+    LLVMOrcLookupStateContinueLookup(LS, Err);
+  };
+  *LookupState = nullptr;
+  return LLVMErrorSuccess;
+}
+
+static void DisposeSuspendedLookupContext(void *Ctx) {
+  static_cast<SuspendedLookupContext *>(Ctx)->Disposed = true;
+}
+
+static void
+suspendLookupTestLookupHandlerCallback(LLVMErrorRef Err,
+                                       LLVMOrcCSymbolMapPairs Result,
+                                       size_t NumPairs, void *RawCtx) {
+  if (Err) {
+    FAIL() << "Suspended DefinitionGenerator did not create symbol \"foo\": "
+           << toString(Err);
+    return;
+  }
+
+  EXPECT_EQ(NumPairs, 1U)
+      << "Unexpected number of result entries: expected 1, got " << NumPairs;
+
+  auto *Ctx = static_cast<SuspendedLookupContext *>(RawCtx);
+  EXPECT_EQ(Result[0].Name, Ctx->NameToGenerate);
+  EXPECT_EQ(Result[0].Sym.Address, Ctx->AddrToGenerate);
+
+  Ctx->QueryCompleted = true;
+}
+
+TEST_F(OrcCAPITestBase, SuspendedLookup) {
+  // Test that we can suspend lookup in a custom generator.
+  SuspendedLookupContext Ctx;
+  Ctx.NameToGenerate = LLVMOrcLLJITMangleAndIntern(Jit, "foo");
+  Ctx.AddrToGenerate = 0x1234;
+
+  // Add generator.
+  LLVMOrcJITDylibAddGenerator(MainDylib,
+                              LLVMOrcCreateCustomCAPIDefinitionGenerator(
+                                  &TryToGenerateWithSuspendedLookup, &Ctx,
+                                  DisposeSuspendedLookupContext));
+
+  // Expect no work to do before the lookup.
+  EXPECT_FALSE(Ctx.AsyncWork) << "Unexpected generator work before lookup";
+
+  // Issue lookup. This should trigger the generator, but generation should
+  // be suspended.
+  LLVMOrcCJITDylibSearchOrderElement SO[] = {
+      {MainDylib, LLVMOrcJITDylibLookupFlagsMatchExportedSymbolsOnly}};
+  LLVMOrcRetainSymbolStringPoolEntry(Ctx.NameToGenerate);
+  LLVMOrcCLookupSetElement LS[] = {
+      {Ctx.NameToGenerate, LLVMOrcSymbolLookupFlagsRequiredSymbol}};
+  LLVMOrcExecutionSessionLookup(ExecutionSession, LLVMOrcLookupKindStatic, SO,
+                                1, LS, 1,
+                                suspendLookupTestLookupHandlerCallback, &Ctx);
+
+  // Expect that we now have generator work to do.
+  EXPECT_TRUE(Ctx.AsyncWork)
+      << "Failed to generator (or failed to suspend generator)";
+
+  // Do the work. This should allow the query to complete.
+  Ctx.AsyncWork();
+
+  // Check that the query completed.
+  EXPECT_TRUE(Ctx.QueryCompleted);
+
+  // Release our local copy of the string.
+  LLVMOrcReleaseSymbolStringPoolEntry(Ctx.NameToGenerate);
+
+  // Explicitly tear down the JIT.
+  LLVMOrcDisposeLLJIT(Jit);
+  Jit = nullptr;
+
+  // Check that the generator context was "destroyed".
+  EXPECT_TRUE(Ctx.Disposed);
+}


        


More information about the llvm-commits mailing list