From 697aa732efc6a47954ba5df99ad10facd0a3371b Mon Sep 17 00:00:00 2001
From: Zoltan Varga <vargaz@gmail.com>
Date: Mon, 23 Jan 2023 13:58:12 -0500
Subject: [PATCH] [mono][aot] Free LLVM module after AOT. (#81014)

This is needed to avoid out of memory errors when multiple assemblies
are AOT-ed during dedup.
---
 src/mono/mono/mini/aot-compiler.c | 12 +++++++++--
 src/mono/mono/mini/mini-llvm.c    | 36 +++++++++++++++++++++++++------
 src/mono/mono/mini/mini-llvm.h    |  1 +
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/src/mono/mono/mini/aot-compiler.c b/src/mono/mono/mini/aot-compiler.c
index 1e4f1799bc251..c93d13556c13b 100644
--- a/src/mono/mono/mini/aot-compiler.c
+++ b/src/mono/mono/mini/aot-compiler.c
@@ -13898,7 +13898,13 @@ aot_opts_free (MonoAotOptions *aot_opts)
 static void
 acfg_free (MonoAotCompile *acfg)
 {
-	mono_img_writer_destroy (acfg->w);
+#ifdef ENABLE_LLVM
+	if (acfg->aot_opts.llvm)
+		mono_llvm_free_aot_module ();
+#endif
+
+	if (acfg->w)
+		mono_img_writer_destroy (acfg->w);
 	for (guint32 i = 0; i < acfg->nmethods; ++i)
 		if (acfg->cfgs [i])
 			mono_destroy_compile (acfg->cfgs [i]);
@@ -14833,9 +14839,11 @@ aot_assembly (MonoAssembly *ass, guint32 jit_opts, MonoAotOptions *aot_options)
 
 	dedup_skip_methods (acfg);
 
-	if (acfg->dedup_collect_only)
+	if (acfg->dedup_collect_only) {
 		/* We only collected methods from this assembly */
+		acfg_free (acfg);
 		return 0;
+	}
 
 	current_acfg = NULL;
 
diff --git a/src/mono/mono/mini/mini-llvm.c b/src/mono/mono/mini/mini-llvm.c
index 9201232fc861d..7b37e1afb1c0a 100644
--- a/src/mono/mono/mini/mini-llvm.c
+++ b/src/mono/mono/mini/mini-llvm.c
@@ -608,6 +608,12 @@ create_address (MonoLLVMModule *module, LLVMValueRef value, LLVMTypeRef type)
 	return res;
 }
 
+static void
+address_free (gpointer addr)
+{
+	g_free (addr);
+}
+
 typedef struct {
 	int32_t size;
 	uint32_t align;
@@ -12996,11 +13002,6 @@ mono_llvm_create_aot_module (MonoAssembly *assembly, const char *global_prefix,
 	gboolean llvm_only = (flags & LLVM_MODULE_FLAG_LLVM_ONLY) ? 1 : 0;
 	gboolean interp = (flags & LLVM_MODULE_FLAG_INTERP) ? 1 : 0;
 
-	/* Delete previous module */
-	g_hash_table_destroy (module->plt_entries);
-	if (module->lmodule)
-		LLVMDisposeModule (module->lmodule);
-
 	memset (module, 0, sizeof (aot_module));
 
 	module->lmodule = LLVMModuleCreateWithName ("aot");
@@ -13019,7 +13020,7 @@ mono_llvm_create_aot_module (MonoAssembly *assembly, const char *global_prefix,
 	module->max_got_offset = initial_got_size;
 	module->context = LLVMGetGlobalContext ();
 	module->cfgs = g_ptr_array_new ();
-	module->aotconst_vars = g_hash_table_new (NULL, NULL);
+	module->aotconst_vars = g_hash_table_new_full (NULL, NULL, NULL, address_free);
 	module->llvm_types = g_hash_table_new (NULL, NULL);
 	module->plt_entries = g_hash_table_new (g_str_hash, g_str_equal);
 	module->plt_entries_ji = g_hash_table_new (NULL, NULL);
@@ -13141,6 +13142,29 @@ mono_llvm_create_aot_module (MonoAssembly *assembly, const char *global_prefix,
 	}
 }
 
+void
+mono_llvm_free_aot_module (void)
+{
+	MonoLLVMModule *module = &aot_module;
+
+	if (module->lmodule)
+		LLVMDisposeModule (module->lmodule);
+
+	g_hash_table_destroy (module->aotconst_vars);
+	g_hash_table_destroy (module->llvm_types);
+	g_hash_table_destroy (module->plt_entries);
+	g_hash_table_destroy (module->plt_entries_ji);
+	g_hash_table_destroy (module->direct_callables);
+	g_hash_table_destroy (module->idx_to_lmethod);
+	g_hash_table_destroy (module->method_to_lmethod);
+	g_hash_table_destroy (module->method_to_call_info);
+	g_hash_table_destroy (module->idx_to_unbox_tramp);
+	g_hash_table_destroy (module->no_method_table_lmethods);
+
+	g_ptr_array_free (module->cfgs, TRUE);
+	g_ptr_array_free (module->callsite_list, TRUE);
+}
+
 void
 mono_llvm_fixup_aot_module (void)
 {
diff --git a/src/mono/mono/mini/mini-llvm.h b/src/mono/mono/mini/mini-llvm.h
index 471d369e27be2..5eaa8d309e48a 100644
--- a/src/mono/mono/mini/mini-llvm.h
+++ b/src/mono/mono/mini/mini-llvm.h
@@ -23,6 +23,7 @@ void mono_llvm_init                     (gboolean enable_jit);
 void mono_llvm_emit_method              (MonoCompile *cfg);
 void mono_llvm_emit_call                (MonoCompile *cfg, MonoCallInst *call);
 void mono_llvm_create_aot_module        (MonoAssembly *assembly, const char *global_prefix, int initial_got_size, LLVMModuleFlags flags);
+void mono_llvm_free_aot_module          (void);
 void mono_llvm_emit_aot_module          (const char *filename, const char *cu_name);
 void mono_llvm_emit_aot_file_info       (MonoAotFileInfo *info, gboolean has_jitted_code);
 gpointer mono_llvm_emit_aot_data        (const char *symbol, guint8 *data, int data_len);