|
3 | 3 | #include "mlir/ExecutionEngine/ExecutionEngine.h" |
4 | 4 | #include "mlir/IR/BuiltinOps.h" |
5 | 5 |
|
| 6 | +#include "llvm/ADT/STLExtras.h" |
| 7 | +#include "llvm/ExecutionEngine/SectionMemoryManager.h" |
6 | 8 | #include "llvm/Support/DynamicLibrary.h" |
| 9 | +#include "llvm/Support/Memory.h" |
| 10 | +#include "llvm/Support/Path.h" |
7 | 11 |
|
8 | 12 | #include <cstdio> |
9 | 13 | #ifdef _WIN32 |
@@ -92,6 +96,63 @@ int loadLibrary(mlir::SmallString<256> &libPath, llvm::StringMap<void *> &export |
92 | 96 | return 0; |
93 | 97 | } |
94 | 98 |
|
| 99 | +// JIT'd globals live in RTDyld-allocated RW data sections (plain VirtualAlloc/mmap |
| 100 | +// memory), which the Boehm GC does not scan: for AOT binaries the loader-mapped |
| 101 | +// data segment is registered as a root set automatically, but in the JIT an object |
| 102 | +// reachable only from a global (e.g. a static class member) is collected on the |
| 103 | +// first GC cycle and its memory recycled, corrupting later accesses. Register every |
| 104 | +// RW data section with the GC as a root range. GC_add_roots/GC_remove_roots are |
| 105 | +// resolved dynamically from the already-loaded TypeScriptRuntime library, so this |
| 106 | +// stays inert when the GC is disabled or not present. |
| 107 | +class GCRootsSectionMemoryMapper : public llvm::SectionMemoryManager::MemoryMapper |
| 108 | +{ |
| 109 | + // AllocationPurpose lives in the enclosing SectionMemoryManager, not in the |
| 110 | + // MemoryMapper base class, so it is not found by unqualified lookup here |
| 111 | + using AllocationPurpose = llvm::SectionMemoryManager::AllocationPurpose; |
| 112 | + using GCRootsFn = void (*)(void *, void *); |
| 113 | + |
| 114 | + public: |
| 115 | + llvm::sys::MemoryBlock allocateMappedMemory(AllocationPurpose purpose, size_t numBytes, |
| 116 | + const llvm::sys::MemoryBlock *const nearBlock, unsigned flags, |
| 117 | + std::error_code &errCode) override |
| 118 | + { |
| 119 | + auto block = llvm::sys::Memory::allocateMappedMemory(numBytes, nearBlock, flags, errCode); |
| 120 | + if (purpose == AllocationPurpose::RWData && block.base() != nullptr) |
| 121 | + { |
| 122 | + if (auto addRoots = reinterpret_cast<GCRootsFn>( |
| 123 | + llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("GC_add_roots"))) |
| 124 | + { |
| 125 | + addRoots(block.base(), static_cast<char *>(block.base()) + block.allocatedSize()); |
| 126 | + gcRootBlocks.push_back(block.base()); |
| 127 | + } |
| 128 | + } |
| 129 | + |
| 130 | + return block; |
| 131 | + } |
| 132 | + |
| 133 | + std::error_code protectMappedMemory(const llvm::sys::MemoryBlock &block, unsigned flags) override |
| 134 | + { |
| 135 | + return llvm::sys::Memory::protectMappedMemory(block, flags); |
| 136 | + } |
| 137 | + |
| 138 | + std::error_code releaseMappedMemory(llvm::sys::MemoryBlock &block) override |
| 139 | + { |
| 140 | + if (llvm::is_contained(gcRootBlocks, block.base())) |
| 141 | + { |
| 142 | + if (auto removeRoots = reinterpret_cast<GCRootsFn>( |
| 143 | + llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("GC_remove_roots"))) |
| 144 | + { |
| 145 | + removeRoots(block.base(), static_cast<char *>(block.base()) + block.allocatedSize()); |
| 146 | + } |
| 147 | + } |
| 148 | + |
| 149 | + return llvm::sys::Memory::releaseMappedMemory(block); |
| 150 | + } |
| 151 | + |
| 152 | + private: |
| 153 | + mlir::SmallVector<void *> gcRootBlocks; |
| 154 | +}; |
| 155 | + |
95 | 156 | int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compileOptions) |
96 | 157 | { |
97 | 158 | // to avoid false positive memory leak reports in release builds |
@@ -135,8 +196,17 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile |
135 | 196 | #define LIB_NAME "lib" |
136 | 197 | #define LIB_EXT "so" |
137 | 198 | #endif |
| 199 | + // Only locate a GC runtime when the user has not already supplied one via |
| 200 | + // --shared-libs: loading a second TypeScriptRuntime copy from a different path |
| 201 | + // (e.g. TSLANG_LIB_PATH) creates two independent GC instances, so roots |
| 202 | + // registered in one are invisible to collections running in the other and |
| 203 | + // GC-heap pointers cross between the two heaps. |
| 204 | + auto hasTypeScriptRuntime = llvm::any_of(clSharedLibs, [](const std::string &libPath) { |
| 205 | + return llvm::sys::path::stem(libPath).contains_insensitive("TypeScriptRuntime"); |
| 206 | + }); |
| 207 | + |
138 | 208 | std::string pathTypeScriptLib("../lib/" LIB_NAME "TypeScriptRuntime." LIB_EXT); |
139 | | - if (!disableGC.getValue()) |
| 209 | + if (!disableGC.getValue() && !hasTypeScriptRuntime) |
140 | 210 | { |
141 | 211 | auto absPath3 = makeAbsolutePath(mergeWithDefaultLibPath(getTslangLibPath(), LIB_NAME "TypeScriptRuntime." LIB_EXT)); |
142 | 212 | if (absPath3.empty()) |
@@ -175,7 +245,11 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile |
175 | 245 | mlir::SmallVector<mlir::StringRef> sharedLibPaths; |
176 | 246 | sharedLibPaths.append(begin(clSharedLibs), end(clSharedLibs)); |
177 | 247 |
|
| 248 | + // must outlive the engine: sections are released from the engine's destructor |
| 249 | + GCRootsSectionMemoryMapper gcRootsMemoryMapper; |
| 250 | + |
178 | 251 | mlir::ExecutionEngineOptions engineOptions; |
| 252 | + engineOptions.sectionMemoryMapper = &gcRootsMemoryMapper; |
179 | 253 | engineOptions.transformer = optPipeline; |
180 | 254 | engineOptions.enableObjectDump = dumpObjectFile; |
181 | 255 | engineOptions.enableGDBNotificationListener = !enableOpt; |
@@ -264,13 +338,29 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile |
264 | 338 | } |
265 | 339 |
|
266 | 340 | // The JIT program has finished. On Windows/COFF the MLIR/ORC LLJIT platform |
267 | | - // registers process-level atexit glue that faults during teardown, so the |
268 | | - // normal CRT exit path crashes after a successful run. Flush stdio and exit |
269 | | - // the process directly, bypassing the CRT atexit handlers. |
| 341 | + // registers process-level atexit glue that faults during teardown (exception |
| 342 | + // 0x80000003), so the normal CRT exit path crashes after a successful run. |
| 343 | + // ExitProcess is not safe either: it runs DLL_PROCESS_DETACH, and |
| 344 | + // TypeScriptRuntime.dll's static AsyncRuntime destructor then waits on its |
| 345 | + // thread pool whose worker threads ExitProcess has already terminated, |
| 346 | + // deadlocking forever when async work is still pending. TerminateProcess |
| 347 | + // skips all teardown, so nothing can crash or deadlock — but it also skips |
| 348 | + // the detach-time stdio flush, so flush every CRT the JIT'd code may have |
| 349 | + // bound to (ucrtbase buffers its streams separately from our static CRT). |
270 | 350 | fflush(stdout); |
271 | 351 | fflush(stderr); |
272 | 352 | #ifdef _WIN32 |
273 | | - ExitProcess(0); |
| 353 | + for (auto crtName : {"ucrtbase.dll", "ucrtbased.dll"}) |
| 354 | + { |
| 355 | + if (HMODULE crt = GetModuleHandleA(crtName)) |
| 356 | + { |
| 357 | + if (auto crtFflush = (int(__cdecl *)(FILE *))GetProcAddress(crt, "fflush")) |
| 358 | + { |
| 359 | + crtFflush(nullptr); // fflush(NULL) flushes all of that CRT's output streams |
| 360 | + } |
| 361 | + } |
| 362 | + } |
| 363 | + TerminateProcess(GetCurrentProcess(), 0); |
274 | 364 | #endif |
275 | 365 | return 0; |
276 | 366 | } |
0 commit comments