Skip to content

Commit 3b991f8

Browse files
danielgustafssondanielgustafsson
committed
llvmjit: Use explicit LLVMContextRef for inlining
When performing inlining LLVM unfortunately "leaks" types (the types survive and are usable, but a new round of inlining will recreate new structurally equivalent types). This accumulation will over time amount to a memory leak which for some queries can be large enough to trigger the OOM process killer. To avoid accumulation of types, all IR related data is stored in an LLVMContextRef which is dropped and recreated in order to release all types. Dropping and recreating incurs overhead, so it will be done only after 100 queries. This is a heuristic which might be revisited, but until we can get the size of the context from LLVM we are flying a bit blind. This issue has been reported several times, there may be more references to it in the archives on top of the threads linked below. This is a backpatch of 9dce220 to all supported branches. Reported-By: Justin Pryzby <pryzby@telsasoft.com> Reported-By: Kurt Roeckx <kurt@roeckx.be> Reported-By: Jaime Casanova <jcasanov@systemguards.com.ec> Reported-By: Lauri Laanmets <pcspets@gmail.com> Author: Andres Freund and Daniel Gustafsson Discussion: https://postgr.es/m/7acc8678-df5f-4923-9cf6-e843131ae89d@www.fastmail.com Discussion: https://postgr.es/m/20201218235607.GC30237@telsasoft.com Discussion: https://postgr.es/m/CAPH-tTxLf44s3CvUUtQpkDr1D8Hxqc2NGDzGXS1ODsfiJ6WSqA@mail.gmail.com Backpatch-through: v12
1 parent c980eed commit 3b991f8

File tree

6 files changed

+252
-109
lines changed

6 files changed

+252
-109
lines changed

‎src/backend/jit/llvm/llvmjit.c

Lines changed: 124 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@
5151
#endif
5252

5353

54+
#define LLVMJIT_LLVM_CONTEXT_REUSE_MAX 100
55+
5456
/* Handle of a module emitted via ORC JIT */
5557
typedef struct LLVMJitHandle
5658
{
@@ -103,8 +105,15 @@ LLVMModuleRef llvm_types_module = NULL;
103105

104106
static bool llvm_session_initialized = false;
105107
static size_t llvm_generation = 0;
108+
109+
/* number of LLVMJitContexts that currently are in use */
110+
static size_t llvm_jit_context_in_use_count = 0;
111+
112+
/* how many times has the current LLVMContextRef been used */
113+
static size_t llvm_llvm_context_reuse_count = 0;
106114
static const char *llvm_triple = NULL;
107115
static const char *llvm_layout = NULL;
116+
static LLVMContextRef llvm_context;
108117

109118

110119
static LLVMTargetRef llvm_targetref;
@@ -125,6 +134,8 @@ static void llvm_compile_module(LLVMJitContext *context);
125134
static void llvm_optimize_module(LLVMJitContext *context, LLVMModuleRef module);
126135

127136
static void llvm_create_types(void);
137+
static void llvm_set_target(void);
138+
static void llvm_recreate_llvm_context(void);
128139
static uint64_t llvm_resolve_symbol(const char *name, void *ctx);
129140

130141
#if LLVM_VERSION_MAJOR > 11
@@ -146,6 +157,63 @@ _PG_jit_provider_init(JitProviderCallbacks *cb)
146157
cb->compile_expr = llvm_compile_expr;
147158
}
148159

160+
161+
/*
162+
* Every now and then create a new LLVMContextRef. Unfortunately, during every
163+
* round of inlining, types may "leak" (they can still be found/used via the
164+
* context, but new types will be created the next time in inlining is
165+
* performed). To prevent that from slowly accumulating problematic amounts of
166+
* memory, recreate the LLVMContextRef we use. We don't want to do so too
167+
* often, as that implies some overhead (particularly re-loading the module
168+
* summaries / modules is fairly expensive). A future TODO would be to make
169+
* this more finegrained and only drop/recreate the LLVMContextRef when we know
170+
* there has been inlining. If we can get the size of the context from LLVM
171+
* then that might be a better way to determine when to drop/recreate rather
172+
* then the usagecount heuristic currently employed.
173+
*/
174+
static void
175+
llvm_recreate_llvm_context(void)
176+
{
177+
if (!llvm_context)
178+
elog(ERROR, "Trying to recreate a non-existing context");
179+
180+
/*
181+
* We can only safely recreate the LLVM context if no other code is being
182+
* JITed, otherwise we'd release the types in use for that.
183+
*/
184+
if (llvm_jit_context_in_use_count > 0)
185+
{
186+
llvm_llvm_context_reuse_count++;
187+
return;
188+
}
189+
190+
if (llvm_llvm_context_reuse_count <= LLVMJIT_LLVM_CONTEXT_REUSE_MAX)
191+
{
192+
llvm_llvm_context_reuse_count++;
193+
return;
194+
}
195+
196+
/*
197+
* Need to reset the modules that the inlining code caches before
198+
* disposing of the context. LLVM modules exist within a specific LLVM
199+
* context, therefore disposing of the context before resetting the cache
200+
* would lead to dangling pointers to modules.
201+
*/
202+
llvm_inline_reset_caches();
203+
204+
LLVMContextDispose(llvm_context);
205+
llvm_context = LLVMContextCreate();
206+
llvm_llvm_context_reuse_count = 0;
207+
208+
/*
209+
* Re-build cached type information, so code generation code can rely on
210+
* that information to be present (also prevents the variables to be
211+
* dangling references).
212+
*/
213+
llvm_create_types();
214+
}
215+
216+
149217
/*
150218
* Create a context for JITing work.
151219
*
@@ -162,6 +230,8 @@ llvm_create_context(int jitFlags)
162230

163231
llvm_session_initialize();
164232

233+
llvm_recreate_llvm_context();
234+
165235
ResourceOwnerEnlargeJIT(CurrentResourceOwner);
166236

167237
context = MemoryContextAllocZero(TopMemoryContext,
@@ -172,6 +242,8 @@ llvm_create_context(int jitFlags)
172242
context->base.resowner = CurrentResourceOwner;
173243
ResourceOwnerRememberJIT(CurrentResourceOwner, PointerGetDatum(context));
174244

245+
llvm_jit_context_in_use_count++;
246+
175247
return context;
176248
}
177249

@@ -181,7 +253,13 @@ llvm_create_context(int jitFlags)
181253
static void
182254
llvm_release_context(JitContext *context)
183255
{
184-
LLVMJitContext *llvm_context = (LLVMJitContext *) context;
256+
LLVMJitContext *llvm_jit_context = (LLVMJitContext *) context;
257+
258+
/*
259+
* Consider as cleaned up even if we skip doing so below, that way we can
260+
* verify the tracking is correct (see llvm_shutdown()).
261+
*/
262+
llvm_jit_context_in_use_count--;
185263

186264
/*
187265
* When this backend is exiting, don't clean up LLVM. As an error might
@@ -193,18 +271,18 @@ llvm_release_context(JitContext *context)
193271

194272
llvm_enter_fatal_on_oom();
195273

196-
if (llvm_context->module)
274+
if (llvm_jit_context->module)
197275
{
198-
LLVMDisposeModule(llvm_context->module);
199-
llvm_context->module = NULL;
276+
LLVMDisposeModule(llvm_jit_context->module);
277+
llvm_jit_context->module = NULL;
200278
}
201279

202-
while (llvm_context->handles != NIL)
280+
while (llvm_jit_context->handles != NIL)
203281
{
204282
LLVMJitHandle *jit_handle;
205283

206-
jit_handle = (LLVMJitHandle *) linitial(llvm_context->handles);
207-
llvm_context->handles = list_delete_first(llvm_context->handles);
284+
jit_handle = (LLVMJitHandle *) linitial(llvm_jit_context->handles);
285+
llvm_jit_context->handles = list_delete_first(llvm_jit_context->handles);
208286

209287
#if LLVM_VERSION_MAJOR > 11
210288
{
@@ -232,6 +310,8 @@ llvm_release_context(JitContext *context)
232310

233311
pfree(jit_handle);
234312
}
313+
list_free(llvm_jit_context->handles);
314+
llvm_jit_context->handles = NIL;
235315

236316
llvm_leave_fatal_on_oom();
237317
}
@@ -251,7 +331,7 @@ llvm_mutable_module(LLVMJitContext *context)
251331
{
252332
context->compiled = false;
253333
context->module_generation = llvm_generation++;
254-
context->module = LLVMModuleCreateWithName("pg");
334+
context->module = LLVMModuleCreateWithNameInContext("pg", llvm_context);
255335
LLVMSetTarget(context->module, llvm_triple);
256336
LLVMSetDataLayout(context->module, llvm_layout);
257337
}
@@ -835,6 +915,14 @@ llvm_session_initialize(void)
835915
LLVMInitializeNativeAsmPrinter();
836916
LLVMInitializeNativeAsmParser();
837917

918+
if (llvm_context == NULL)
919+
{
920+
llvm_context = LLVMContextCreate();
921+
922+
llvm_jit_context_in_use_count = 0;
923+
llvm_llvm_context_reuse_count = 0;
924+
}
925+
838926
/*
839927
* When targeting LLVM 15, turn off opaque pointers for the context we
840928
* build our code in. We don't need to do so for other contexts (e.g.
@@ -854,6 +942,11 @@ llvm_session_initialize(void)
854942
*/
855943
llvm_create_types();
856944

945+
/*
946+
* Extract target information from loaded module.
947+
*/
948+
llvm_set_target();
949+
857950
if (LLVMGetTargetFromTriple(llvm_triple, &llvm_targetref, &error) != 0)
858951
{
859952
elog(FATAL, "failed to query triple %s", error);
@@ -949,6 +1042,10 @@ llvm_shutdown(int code, Datum arg)
9491042
return;
9501043
}
9511044

1045+
if (llvm_jit_context_in_use_count != 0)
1046+
elog(PANIC, "LLVMJitContext in use count not 0 at exit (is %zu)",
1047+
llvm_jit_context_in_use_count);
1048+
9521049
#if LLVM_VERSION_MAJOR > 11
9531050
{
9541051
if (llvm_opt3_orc)
@@ -1011,6 +1108,23 @@ load_return_type(LLVMModuleRef mod, const char *name)
10111108
return typ;
10121109
}
10131110

1111+
/*
1112+
* Load triple & layout from clang emitted file so we're guaranteed to be
1113+
* compatible.
1114+
*/
1115+
static void
1116+
llvm_set_target(void)
1117+
{
1118+
if (!llvm_types_module)
1119+
elog(ERROR, "failed to extract target information, llvmjit_types.c not loaded");
1120+
1121+
if (llvm_triple == NULL)
1122+
llvm_triple = pstrdup(LLVMGetTarget(llvm_types_module));
1123+
1124+
if (llvm_layout == NULL)
1125+
llvm_layout = pstrdup(LLVMGetDataLayoutStr(llvm_types_module));
1126+
}
1127+
10141128
/*
10151129
* Load required information, types, function signatures from llvmjit_types.c
10161130
* and make them available in global variables.
@@ -1034,19 +1148,12 @@ llvm_create_types(void)
10341148
}
10351149

10361150
/* eagerly load contents, going to need it all */
1037-
if (LLVMParseBitcode2(buf, &llvm_types_module))
1151+
if (LLVMParseBitcodeInContext2(llvm_context, buf, &llvm_types_module))
10381152
{
1039-
elog(ERROR, "LLVMParseBitcode2 of %s failed", path);
1153+
elog(ERROR, "LLVMParseBitcodeInContext2 of %s failed", path);
10401154
}
10411155
LLVMDisposeMemoryBuffer(buf);
10421156

1043-
/*
1044-
* Load triple & layout from clang emitted file so we're guaranteed to be
1045-
* compatible.
1046-
*/
1047-
llvm_triple = pstrdup(LLVMGetTarget(llvm_types_module));
1048-
llvm_layout = pstrdup(LLVMGetDataLayoutStr(llvm_types_module));
1049-
10501157
TypeSizeT = llvm_pg_var_type("TypeSizeT");
10511158
TypeParamBool = load_return_type(llvm_types_module, "FunctionReturningBool");
10521159
TypeStorageBool = llvm_pg_var_type("TypeStorageBool");

0 commit comments

Comments
 (0)