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

[sunho]

Summary

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.

[sgraenitz]

Thanks, that sounds reasonable. Maybe we can provide some more context on atexit vs. __cxa_atexit in the commit message?

It appears that there is a test for __cxa_atexit here -- can we extend it to cover atexit?
https://github.com/llvm/llvm-project/blob/main/llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll

[sunho]

Add testcases

[sunho]

@sgraenitz I added a testcase and editted the description of the revision to explain the situation more thoroughly which I will use as a commit message. Does it look good to you?

[sgraenitz]

Thanks, yes looks good to me!

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

setInitTransform(J, GlobalCtorDtorScraper(*this, InitFunctionPrefix,
                                               DeInitFunctionPrefix));
 
     SymbolMap StdInterposes;
 
     StdInterposes[J.mangleAndIntern("__lljit.platform_support_instance")] =
         JITEvaluatedSymbol(pointerToJITTargetAddress(this),
                            JITSymbolFlags::Exported);
     StdInterposes[J.mangleAndIntern("__lljit.cxa_atexit_helper")] =
-        JITEvaluatedSymbol(pointerToJITTargetAddress(registerAtExitHelper),
+        JITEvaluatedSymbol(pointerToJITTargetAddress(registerCxaAtExitHelper),
                            JITSymbolFlags());
 
     cantFail(
         J.getMainJITDylib().define(absoluteSymbols(std::move(StdInterposes))));
     cantFail(setupJITDylib(J.getMainJITDylib()));
     cantFail(J.addIRModule(J.getMainJITDylib(), createPlatformRuntimeModule()));
   }
 
   ExecutionSession &getExecutionSession() { return J.getExecutionSession(); }
 
   /// Adds a module that defines the __dso_handle global.
   Error setupJITDylib(JITDylib &JD) {
 
     // Add per-jitdylib standard interposes.
     SymbolMap PerJDInterposes;
     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>();
     auto M = std::make_unique<Module>("__standard_lib", *Ctx);
     M->setDataLayout(J.getDataLayout());
 
     auto *Int64Ty = Type::getInt64Ty(*Ctx);
     auto *DSOHandle = new GlobalVariable(
auto *PlatformInstanceDecl = new GlobalVariable(
         nullptr, "__lljit.platform_support_instance");
 
     auto *VoidTy = Type::getVoidTy(*Ctx);
     addHelperAndWrapper(
         *M, "__lljit_run_atexits", FunctionType::get(VoidTy, {}, false),
         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)));
   }
 
   Error notifyAdding(ResourceTracker &RT, const MaterializationUnit &MU) {
     auto &JD = RT.getJITDylib();
     if (auto &InitSym = MU.getInitializerSymbol())
       InitSymbols[&JD].add(InitSym, SymbolLookupFlags::WeaklyReferencedSymbol);
     else {
if (auto Err = getExecutionSession().runSessionLocked([&]() -> Error {
         }))
       return Err;
 
     return Platform::lookupInitSymbols(getExecutionSession(),
                                        RequiredInitSymbols)
         .takeError();
   }
 
-  static void registerAtExitHelper(void *Self, void (*F)(void *), void *Ctx,
+  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 "
              << (*static_cast<JITDylib **>(DSOHandle))->getName() << "\n";
     });
     static_cast<GenericLLVMIRPlatformSupport *>(Self)->AtExitMgr.runAtExits(
         DSOHandle);
   }
ThreadSafeModule createPlatformRuntimeModule() {
     auto *PlatformInstanceDecl = new GlobalVariable(
         *M, GenericIRPlatformSupportTy, true, GlobalValue::ExternalLinkage,
         nullptr, "__lljit.platform_support_instance");
 
     auto *Int8Ty = Type::getInt8Ty(*Ctx);
     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});
 
     return ThreadSafeModule(std::move(M), std::move(Ctx));
   }
 
   LLJIT &J;

llvm/test/ExecutionEngine/OrcLazy/global-ctors-and-dtors.ll

 ; Test that global constructors and destructors are run:
 ;
 ; RUN: lli -jit-kind=orc-lazy -orc-lazy-debug=funcs-to-stdout -extra-module %s \
 ; RUN:   %S/Inputs/noop-main.ll | FileCheck %s
 ;
 ; Test that this is true for global constructors and destructors in other
 ; JITDylibs.
 ; RUN: lli -jit-kind=orc-lazy -orc-lazy-debug=funcs-to-stdout \
 ; RUN:   -jd extra -extra-module %s -jd main %S/Inputs/noop-main.ll | FileCheck %s
 ;
 ; 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 }
 
 @f = global %class.Foo zeroinitializer, align 1
 @__dso_handle = external global i8
 @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
 }
 
 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
 }