[ORC][LLJIT] Define atexit symbol in GenericLLVMIRPlatformSupport.

Define atexit symbol in GenericLLVMIRPlatformSupport so that it doesn't need to be defined by user.

On windows, llvm codegen emits atexit runtime calls to support global deinitializers as there is no lower function like cxa_atexit as in Itanium C++ ABI. ORC JIT user had to define custom atexit symbol manually. This was a hassle as it has to deal with dso_handle and cxa_atexit internals of LLJIT. If client didn't provide atexit definition, the default behaviour is just linking with host atexit function which is destined to fail as it calls dtors when the host program exits. This is after jit instances and buffers are freed, so users would see weird access violation exception from the uknown location. (in console application, the debugger thinks exception happened in scrt_common_main_seh)

This is a hack that has some caveats. (e.g. memory address is not identical) But, it's better than the situation described in the above. Ultimately, we will move on to ORC runtime that is able to solve the memory address issue properly.

Reviewed By: sgraenitz

Differential Revision: https://reviews.llvm.org/D128037

2 files changed, 47 insertions, 15 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index c6442e1..6d67e6d 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -143,7 +143,7 @@ public:
         JITEvaluatedSymbol(pointerToJITTargetAddress(this),
                            JITSymbolFlags::Exported);
     StdInterposes[J.mangleAndIntern("__lljit.cxa_atexit_helper")] =
-        JITEvaluatedSymbol(pointerToJITTargetAddress(registerAtExitHelper),
+        JITEvaluatedSymbol(pointerToJITTargetAddress(registerCxaAtExitHelper),
                            JITSymbolFlags());
 
     cantFail(
@@ -162,6 +162,9 @@ public:
     PerJDInterposes[J.mangleAndIntern("__lljit.run_atexits_helper")] =
         JITEvaluatedSymbol(pointerToJITTargetAddress(runAtExitsHelper),
                            JITSymbolFlags());
+    PerJDInterposes[J.mangleAndIntern("__lljit.atexit_helper")] =
+        JITEvaluatedSymbol(pointerToJITTargetAddress(registerAtExitHelper),
+                           JITSymbolFlags());
     cantFail(JD.define(absoluteSymbols(std::move(PerJDInterposes))));
 
     auto Ctx = std::make_unique<LLVMContext>();
@@ -190,6 +193,14 @@ public:
         GlobalValue::HiddenVisibility, "__lljit.run_atexits_helper",
         {PlatformInstanceDecl, DSOHandle});
 
+    auto *IntTy = Type::getIntNTy(*Ctx, sizeof(int) * CHAR_BIT);
+    auto *AtExitCallbackTy = FunctionType::get(VoidTy, {}, false);
+    auto *AtExitCallbackPtrTy = PointerType::getUnqual(AtExitCallbackTy);
+    addHelperAndWrapper(*M, "atexit",
+                        FunctionType::get(IntTy, {AtExitCallbackPtrTy}, false),
+                        GlobalValue::HiddenVisibility, "__lljit.atexit_helper",
+                        {PlatformInstanceDecl, DSOHandle});
+
     return J.addIRModule(JD, ThreadSafeModule(std::move(M), std::move(Ctx)));
   }
 
@@ -413,16 +424,25 @@ private:
         .takeError();
   }
 
-  static void registerAtExitHelper(void *Self, void (*F)(void *), void *Ctx,
-                                   void *DSOHandle) {
+  static void registerCxaAtExitHelper(void *Self, void (*F)(void *), void *Ctx,
+                                      void *DSOHandle) {
     LLVM_DEBUG({
-      dbgs() << "Registering atexit function " << (void *)F << " for JD "
+      dbgs() << "Registering cxa atexit function " << (void *)F << " for JD "
              << (*static_cast<JITDylib **>(DSOHandle))->getName() << "\n";
     });
     static_cast<GenericLLVMIRPlatformSupport *>(Self)->AtExitMgr.registerAtExit(
         F, Ctx, DSOHandle);
   }
 
+  static void registerAtExitHelper(void *Self, void *DSOHandle, void (*F)()) {
+    LLVM_DEBUG({
+      dbgs() << "Registering atexit function " << (void *)F << " for JD "
+             << (*static_cast<JITDylib **>(DSOHandle))->getName() << "\n";
+    });
+    static_cast<GenericLLVMIRPlatformSupport *>(Self)->AtExitMgr.registerAtExit(
+        reinterpret_cast<void (*)(void *)>(F), nullptr, DSOHandle);
+  }
+
   static void runAtExitsHelper(void *Self, void *DSOHandle) {
     LLVM_DEBUG({
       dbgs() << "Running atexit functions for JD "
@@ -450,12 +470,12 @@ private:
     auto *IntTy = Type::getIntNTy(*Ctx, sizeof(int) * CHAR_BIT);
     auto *VoidTy = Type::getVoidTy(*Ctx);
     auto *BytePtrTy = PointerType::getUnqual(Int8Ty);
-    auto *AtExitCallbackTy = FunctionType::get(VoidTy, {BytePtrTy}, false);
-    auto *AtExitCallbackPtrTy = PointerType::getUnqual(AtExitCallbackTy);
+    auto *CxaAtExitCallbackTy = FunctionType::get(VoidTy, {BytePtrTy}, false);
+    auto *CxaAtExitCallbackPtrTy = PointerType::getUnqual(CxaAtExitCallbackTy);
 
     addHelperAndWrapper(
         *M, "__cxa_atexit",
-        FunctionType::get(IntTy, {AtExitCallbackPtrTy, BytePtrTy, BytePtrTy},
+        FunctionType::get(IntTy, {CxaAtExitCallbackPtrTy, BytePtrTy, BytePtrTy},
                           false),
         GlobalValue::DefaultVisibility, "__lljit.cxa_atexit_helper",
         {PlatformInstanceDecl});
diff --git a/llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll b/llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll
index 961f9aa..286e136 100644
--- a/llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll
+++ b/llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll
@@ -11,8 +11,9 @@
 ; CHECK: Hello from constructor
 ; CHECK: Hello
 ; CHECK: [ {{.*}}main{{.*}} ]
-; CHECK: Goodbye
-; CHECK: Goodbye again
+; CHECK: Goodbye from atexit
+; CHECK: Goodbye from __cxa_atexit
+; CHECK: Goodbye from destructor
 
 %class.Foo = type { i8 }
 
@@ -21,28 +22,39 @@
 @llvm.global_ctors = appending global [2 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_hello.cpp, i8* null }, { i32, void ()*, i8* } { i32 1024, void ()* @constructor, i8* null }]
 @llvm.global_dtors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 0, void ()* @printf_wrapper, i8* null }]
 @str = private unnamed_addr constant [6 x i8] c"Hello\00"
-@str2 = private unnamed_addr constant [8 x i8] c"Goodbye\00"
-@str3 = global [14 x i8] c"Goodbye again\00"
-@str4 = private unnamed_addr constant [23 x i8] c"Hello from constructor\00"
+@str2 = private unnamed_addr constant [23 x i8] c"Hello from constructor\00"
+@str3 = private unnamed_addr constant [24 x i8] c"Goodbye from destructor\00"
+@str4 = global [26 x i8] c"Goodbye from __cxa_atexit\00"
+@str5 = global [20 x i8] c"Goodbye from atexit\00"
+
 
 define linkonce_odr void @_ZN3FooD1Ev(%class.Foo* nocapture readnone %this) unnamed_addr align 2 {
 entry:
-  %puts.i = tail call i32 @puts(i8* getelementptr inbounds ([8 x i8], [8 x i8]* @str2, i64 0, i64 0))
+  %puts.i = tail call i32 @puts(i8* getelementptr inbounds ([26 x i8], [26 x i8]* @str4, i64 0, i64 0))
+  ret void
+}
+
+define void @atexit_handler() {
+entry:
+  %puts.i = tail call i32 @puts(i8* getelementptr inbounds ([20 x i8], [20 x i8]* @str5, i64 0, i64 0))
   ret void
 }
 
 declare i32 @__cxa_atexit(void (i8*)*, i8*, i8*)
 
+declare i32 @atexit(void ()*)
+
 define internal void @_GLOBAL__sub_I_hello.cpp() {
 entry:
   %puts.i.i.i = tail call i32 @puts(i8* getelementptr inbounds ([6 x i8], [6 x i8]* @str, i64 0, i64 0))
   %0 = tail call i32 @__cxa_atexit(void (i8*)* bitcast (void (%class.Foo*)* @_ZN3FooD1Ev to void (i8*)*), i8* getelementptr inbounds (%class.Foo, %class.Foo* @f, i64 0, i32 0), i8* @__dso_handle)
+  %1 = tail call i32 @atexit(void ()* @atexit_handler)
   ret void
 }
 
 define void @printf_wrapper() {
 entry:
-  %0 = tail call i32 @puts(i8* getelementptr inbounds ([14 x i8], [14 x i8]* @str3, i64 0, i64 0))
+  %0 = tail call i32 @puts(i8* getelementptr inbounds ([24 x i8], [24 x i8]* @str3, i64 0, i64 0))
   ret void
 }
 
@@ -50,6 +62,6 @@ declare i32 @puts(i8* nocapture readonly)
 
 define void @constructor() {
 entry:
-  %0 = tail call i32 @puts(i8* getelementptr inbounds ([23 x i8], [23 x i8]* @str4, i64 0, i64 0))
+  %0 = tail call i32 @puts(i8* getelementptr inbounds ([23 x i8], [23 x i8]* @str2, i64 0, i64 0))
   ret void
 }