From 951a39f0722039bfc22f7a89b3e84c68b6c7abfa Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 16:26:51 +0000 Subject: [PATCH 1/2] JIT cache: each .ts file is compiled into an object kept in __jit Under --emit=jit the program and every .ts module it imports (directly or through another) are each compiled into an object of their own, stored in a `__jit` folder beside the source with a manifest of what it was compiled from (the sources read - itself, references, imports, lib.d.ts - the imported shared libraries, and the tslang build). The next run loads the objects into the LLJIT instead of compiling again while all of that is unchanged; an edited module is recompiled together with its importers. Options are part of the object's name, so runs with other options keep objects of their own. - An import pointing to a .ts file is a declaration again under the JIT, as it is compiled; the module's object provides the bodies. MLIRGen records the .ts files imported (dependencies first) and the libraries imported. - Each object's llvm.global_ctors are moved into one init function, called in dependency order before the program starts (the JIT runs no object's global_ctors). On ELF, COMDAT definitions (a class's `.size`, defined by the module and its importers) are made weak so ORC keeps one. - A program in an import cycle is compiled as one module with the imports' bodies included (a module in a cycle can't be compiled on its own, the same holds for --emit=obj) and cached as one object. - Writes are atomic (temp file + rename), so concurrent runs share a cache. - --jit-cache=false keeps the previous behaviour (one module, every run); --jit-cache-dir= puts all objects in one folder. Not used for stdin input or --dump-object-file. - runJit is split into reusable parts (process setup, library loading, JIT creation, program run) shared with the cached path; on Win64 the _CxxThrowException shim now finds the image base of the object the ThrowInfo is in, since there can be more than one JIT'd object. Tests: jit-cache/ (diamond imports with initialization order, an import cycle) and jit-cache.cmake (cache reused without rewriting, recompiled after an imported module changes). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VPAHhg8c9NEdoYNw52AtTg --- .gitignore | 3 + README.md | 20 + tslang/include/TypeScript/DataStructs.h | 15 + tslang/lib/TypeScript/MLIRGenModule.cpp | 34 +- tslang/test/tester/CMakeLists.txt | 14 + tslang/test/tester/jit-cache.cmake | 76 +++ tslang/test/tester/jit-cache/base_module.ts | 23 + tslang/test/tester/jit-cache/cycle_main.ts | 11 + tslang/test/tester/jit-cache/cycle_module.ts | 5 + tslang/test/tester/jit-cache/greeting.ts | 3 + .../test/tester/jit-cache/greeting_module.ts | 3 + .../test/tester/jit-cache/imports_diamond.ts | 25 + tslang/test/tester/jit-cache/other_module.ts | 5 + .../tester/jit-cache/sub/derived_module.ts | 14 + tslang/tslang/CMakeLists.txt | 1 + tslang/tslang/jit.cpp | 483 +++++++++---- tslang/tslang/jitcache.cpp | 637 ++++++++++++++++++ tslang/tslang/jitcache.h | 34 + tslang/tslang/opts.cpp | 5 + tslang/tslang/tslang.cpp | 60 +- 20 files changed, 1295 insertions(+), 176 deletions(-) create mode 100644 tslang/test/tester/jit-cache.cmake create mode 100644 tslang/test/tester/jit-cache/base_module.ts create mode 100644 tslang/test/tester/jit-cache/cycle_main.ts create mode 100644 tslang/test/tester/jit-cache/cycle_module.ts create mode 100644 tslang/test/tester/jit-cache/greeting.ts create mode 100644 tslang/test/tester/jit-cache/greeting_module.ts create mode 100644 tslang/test/tester/jit-cache/imports_diamond.ts create mode 100644 tslang/test/tester/jit-cache/other_module.ts create mode 100644 tslang/test/tester/jit-cache/sub/derived_module.ts create mode 100644 tslang/tslang/jitcache.cpp create mode 100644 tslang/tslang/jitcache.h diff --git a/.gitignore b/.gitignore index e8a502f7c..b1394d0c0 100644 --- a/.gitignore +++ b/.gitignore @@ -12,3 +12,6 @@ package-lock.json node_modules gc-* libatomic_ops-* + +# the JIT cache (tslang --emit=jit) +__jit/ diff --git a/README.md b/README.md index 6f632bb6b..e5cb8e490 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,26 @@ Result Hello World! ``` +### JIT cache + +The JIT compiles the program and each `.ts` module it imports into an object file of its own and +keeps it in a `__jit` folder next to the source file. On the next run it loads those objects +instead of compiling the files again, as long as nothing they were compiled from has changed: the +file itself, the files it references or imports, `lib.d.ts`, the options, or the `tslang` build. +An edited module is recompiled together with every file that imports it. + +```text +hello.ts +__jit/hello.ts..o the object +__jit/hello.ts..o.deps what it was compiled from +``` + +- `--jit-cache-dir=` keeps all the objects in one folder instead; +- `--jit-cache=false` turns the cache off: everything is compiled into one module on every run. + +A program whose imports form a cycle is compiled into one object. Its modules can't be compiled +one by one, which is also true for `--emit=obj`. + ## Debugging JIT code with GDB (Linux) JIT-compiled TypeScript can be debugged at source level with GDB — breakpoints on `.ts` lines, diff --git a/tslang/include/TypeScript/DataStructs.h b/tslang/include/TypeScript/DataStructs.h index 3204ee881..7801e0abf 100644 --- a/tslang/include/TypeScript/DataStructs.h +++ b/tslang/include/TypeScript/DataStructs.h @@ -54,6 +54,21 @@ struct CompileOptions // See docs/single-gc-collector-design.md. bool importsSharedLibrary = false; + // Under the JIT: every .ts module is compiled into an object of its own, kept in a cache + // folder (`__jit` beside the source), and the objects are loaded side by side - so an import + // pointing to a .ts file is a declaration, as it is compiled. Off (`--jit-cache=false`), the + // imported file is included with its bodies into the one module the JIT runs. + bool jitCache = false; + + // Filled while the module is generated, for the JIT cache: the .ts files imported, directly + // or through another import, each after the ones it imports (canonical paths), and the + // shared libraries imported, whose declarations the module was generated against. + std::vector sourceImports; + std::vector sharedLibraryImports; + // an import led back to a file still being generated: the modules in the cycle cannot be + // compiled one by one (the JIT cache then compiles the program as one module) + bool importCycle = false; + // Whether the Boehm runtime has to be present. Only `gc` needs it: it is the model whose // reclamation *is* the collector. `rc` frees through the reference counts it maintains and // `none` frees nothing, so both allocate straight from `malloc` and neither links libgc. diff --git a/tslang/lib/TypeScript/MLIRGenModule.cpp b/tslang/lib/TypeScript/MLIRGenModule.cpp index 433fec27f..09854b044 100644 --- a/tslang/lib/TypeScript/MLIRGenModule.cpp +++ b/tslang/lib/TypeScript/MLIRGenModule.cpp @@ -1068,8 +1068,14 @@ namespace mlirgen // and through another module - so its declarations are all there. Or an import cycle back // to a file still being generated: generating it again would never end; what the cycle // needs from it and it has not declared yet stays unresolved. - if (emittedFiles.contains(canonicalPath) || filesInProgress.contains(canonicalPath)) + if (emittedFiles.contains(canonicalPath)) + { + return mlir::success(); + } + + if (filesInProgress.contains(canonicalPath)) { + compileOptions.importCycle = true; return mlir::success(); } @@ -1087,12 +1093,13 @@ namespace mlirgen auto inProgress = llvm::make_scope_exit([&]() { filesInProgress.erase(canonicalPath); }); // An import that points to a .ts file: compiled, it is a declaration - the module is - // compiled into an object of its own, and the objects are linked into one program. Under - // the JIT there is no other object, so it is included with its bodies, as a referenced - // file is. (An import that points to a library is mlirGenImportSharedLib.) A .d.ts is - // read as declarations either way. + // compiled into an object of its own, and the objects are linked into one program. The + // JIT does the same with its cache (the module's object is loaded beside the program's); + // without it there is no other object, so it is included with its bodies, as a + // referenced file is. (An import that points to a library is mlirGenImportSharedLib.) A + // .d.ts is read as declarations either way. MLIRValueGuard vg(declarationMode); - declarationMode = !compileOptions.isJit; + declarationMode = !compileOptions.isJit || compileOptions.jitCache; // What the imported module exports is its own to export: it is compiled separately, with // its declarations in its own __decls. Added to this module's too, a library holding both @@ -1156,6 +1163,16 @@ namespace mlirgen emittedFiles.insert(canonicalPath); emittedFiles.insert(declSymbol); emittedFiles.insert(genericDeclSymbol); + + // the object the JIT has to load for it; after the ones it imports, which were + // generated inside it. A pass generated again (discovery, then the module) finds it + // listed already. + if (!StringRef(canonicalPath).ends_with(".d.ts") && + !llvm::is_contained(compileOptions.sourceImports, canonicalPath)) + { + compileOptions.sourceImports.push_back(canonicalPath); + } + return mlir::success(); } @@ -1439,6 +1456,11 @@ namespace mlirgen // only now: an import that failed is tried again on the next pass, and must fail again // rather than find itself already done emittedFiles.insert(canonicalPath); + if (!llvm::is_contained(compileOptions.sharedLibraryImports, canonicalPath)) + { + compileOptions.sharedLibraryImports.push_back(canonicalPath); + } + for (auto declSymbol : declaredSymbols) { emittedFiles.insert(declSymbol); diff --git a/tslang/test/tester/CMakeLists.txt b/tslang/test/tester/CMakeLists.txt index f108399aa..8fb9333d9 100644 --- a/tslang/test/tester/CMakeLists.txt +++ b/tslang/test/tester/CMakeLists.txt @@ -2629,6 +2629,20 @@ set(import_source_dir "${PROJECT_SOURCE_DIR}/test/tester/import-source") tslang_add_test(NAME test-jit-import-source-new-imported-class COMMAND test-runner -jit "${import_source_dir}/new_imported_class.ts") tslang_add_import_tests(import-source-new-imported-class "${import_source_dir}" new_imported_class imported_class_module) +# The JIT cache (--emit=jit, on by default): each .ts module imported is an object of its own, +# loaded beside the program's - imported through two others and initialized first; and a program +# in an import cycle, compiled as one module. jit-cache.cmake: the cache is used, and compiled +# again when a source changes. +set(jit_cache_dir "${PROJECT_SOURCE_DIR}/test/tester/jit-cache") +tslang_add_test(NAME test-jit-cache-imports-diamond COMMAND test-runner -jit "${jit_cache_dir}/imports_diamond.ts") +tslang_add_test(NAME test-jit-cache-import-cycle COMMAND test-runner -jit "${jit_cache_dir}/cycle_main.ts") +add_test(NAME test-jit-cache-reuse-and-invalidate + COMMAND ${CMAKE_COMMAND} + "-DTSLANG=$" + "-DSOURCE_DIR=${jit_cache_dir}" + "-DWORK_DIR=${CMAKE_CURRENT_BINARY_DIR}/jit-cache-reuse" + -P "${CMAKE_CURRENT_SOURCE_DIR}/jit-cache.cmake") + # `import { a as b }` and `import * as M` - each used to be "can't resolve name". set(import_bindings_dir "${PROJECT_SOURCE_DIR}/test/tester/import-bindings") tslang_add_import_tests(import-bindings-renamed "${import_bindings_dir}" import_renamed names_module) diff --git a/tslang/test/tester/jit-cache.cmake b/tslang/test/tester/jit-cache.cmake new file mode 100644 index 000000000..28de40a4b --- /dev/null +++ b/tslang/test/tester/jit-cache.cmake @@ -0,0 +1,76 @@ +# The JIT cache: `tslang --emit=jit` compiles the program and each .ts module it imports into an +# object of its own, kept in `__jit` beside the source, and loads the objects from there while +# nothing they were compiled from changes. +# +# 1. The first run compiles both files and leaves an object and a manifest for each in __jit. +# 2. The second run finds them: the program runs, and neither object is written again. +# 3. The imported module changes: the program runs the new code, and both objects are compiled +# again - the program's too, which was compiled against the module's declarations. +# +# The sources are copied into the working directory, so the cache and the edit stay out of the +# source tree. -mm=none and --no-default-lib keep the run to the compiler alone. + +cmake_minimum_required(VERSION 3.17.3) + +foreach(var TSLANG SOURCE_DIR WORK_DIR) + if(NOT DEFINED ${var}) + message(FATAL_ERROR "${var} is required") + endif() +endforeach() + +file(REMOVE_RECURSE "${WORK_DIR}") +file(MAKE_DIRECTORY "${WORK_DIR}") +file(COPY "${SOURCE_DIR}/greeting.ts" "${SOURCE_DIR}/greeting_module.ts" DESTINATION "${WORK_DIR}") + +# run() - fails the test unless the program runs and prints what is expected +function(run expected) + execute_process(COMMAND "${TSLANG}" --emit=jit -mm=none --no-default-lib greeting.ts + WORKING_DIRECTORY "${WORK_DIR}" + OUTPUT_VARIABLE out + ERROR_VARIABLE err + RESULT_VARIABLE status) + string(STRIP "${out}" out) + if(NOT status EQUAL 0 OR NOT out STREQUAL expected) + message(FATAL_ERROR "expected '${expected}': exit ${status}\n${out}\n${err}") + endif() +endfunction() + +# the objects in the cache, and when each was written +function(cached_objects result) + file(GLOB objects "${WORK_DIR}/__jit/*.o" "${WORK_DIR}/__jit/*.obj") + list(SORT objects) + set(stamps "") + foreach(object ${objects}) + file(TIMESTAMP "${object}" stamp "%Y%m%d%H%M%S") + list(APPEND stamps "${object}@${stamp}") + endforeach() + set(${result} "${stamps}" PARENT_SCOPE) +endfunction() + +run("one") +cached_objects(first) +list(LENGTH first count) +if(NOT count EQUAL 2) + message(FATAL_ERROR "expected an object for greeting.ts and one for greeting_module.ts in __jit, found: ${first}") +endif() + +# a second apart, so an object written again has another time +execute_process(COMMAND "${CMAKE_COMMAND}" -E sleep 1.1) + +run("one") +cached_objects(second) +if(NOT first STREQUAL second) + message(FATAL_ERROR "the second run compiled again instead of using the cache:\n${first}\n${second}") +endif() + +file(WRITE "${WORK_DIR}/greeting_module.ts" "export function greeting() {\n return \"two\";\n}\n") + +run("two") +cached_objects(third) +foreach(object ${third}) + if(object IN_LIST second) + message(FATAL_ERROR "${object} was not compiled again after greeting_module.ts changed") + endif() +endforeach() + +message(STATUS "the JIT cache is used while the sources are the same, and compiled again when they change") diff --git a/tslang/test/tester/jit-cache/base_module.ts b/tslang/test/tester/jit-cache/base_module.ts new file mode 100644 index 000000000..c598c519f --- /dev/null +++ b/tslang/test/tester/jit-cache/base_module.ts @@ -0,0 +1,23 @@ +function makeStart() { + return 10 + 5; +} + +// not a constant: set by the module's own initializer, which has to run before its importers' +export let counter = makeStart(); + +export function bump() { + counter++; + return counter; +} + +export class Base { + static created = 0; + + constructor() { + Base.created++; + } + + hello() { + return "base"; + } +} diff --git a/tslang/test/tester/jit-cache/cycle_main.ts b/tslang/test/tester/jit-cache/cycle_main.ts new file mode 100644 index 000000000..7861d0f6d --- /dev/null +++ b/tslang/test/tester/jit-cache/cycle_main.ts @@ -0,0 +1,11 @@ +// An import cycle with code at the top level: cycle_module cannot be compiled by itself, so the +// JIT cache compiles the program as one module. +import './cycle_module' + +export function fromMain() { + return "main"; +} + +assert(fromModule() == "module+main"); + +print("done."); diff --git a/tslang/test/tester/jit-cache/cycle_module.ts b/tslang/test/tester/jit-cache/cycle_module.ts new file mode 100644 index 000000000..5afb3b2a7 --- /dev/null +++ b/tslang/test/tester/jit-cache/cycle_module.ts @@ -0,0 +1,5 @@ +import './cycle_main' + +export function fromModule() { + return "module+" + fromMain(); +} diff --git a/tslang/test/tester/jit-cache/greeting.ts b/tslang/test/tester/jit-cache/greeting.ts new file mode 100644 index 000000000..96f9f77e1 --- /dev/null +++ b/tslang/test/tester/jit-cache/greeting.ts @@ -0,0 +1,3 @@ +import './greeting_module' + +print(greeting()); diff --git a/tslang/test/tester/jit-cache/greeting_module.ts b/tslang/test/tester/jit-cache/greeting_module.ts new file mode 100644 index 000000000..1f5457d17 --- /dev/null +++ b/tslang/test/tester/jit-cache/greeting_module.ts @@ -0,0 +1,3 @@ +export function greeting() { + return "one"; +} diff --git a/tslang/test/tester/jit-cache/imports_diamond.ts b/tslang/test/tester/jit-cache/imports_diamond.ts new file mode 100644 index 000000000..abe8fafde --- /dev/null +++ b/tslang/test/tester/jit-cache/imports_diamond.ts @@ -0,0 +1,25 @@ +// Under the JIT cache each module is an object of its own: this program, sub/derived_module and +// other_module, which both import base_module - loaded once, and initialized before either. +import './sub/derived_module' +import './other_module' + +assert(doubled == 30, "base_module initialized before derived_module"); + +const d = new Derived(); +assert(d.hello() == "derived:base"); +assert(Base.created == 1); + +assert(bumpFromDerived() == 16); +assert(bumpFromOther() == 17); +assert(counter == 17, "one base_module, shared by both"); + +let caught = ""; +try { + throw "boom"; +} catch (e) { + caught = e; +} + +assert(caught == "boom"); + +print("done."); diff --git a/tslang/test/tester/jit-cache/other_module.ts b/tslang/test/tester/jit-cache/other_module.ts new file mode 100644 index 000000000..bd1b9ddd2 --- /dev/null +++ b/tslang/test/tester/jit-cache/other_module.ts @@ -0,0 +1,5 @@ +import './base_module' + +export function bumpFromOther() { + return bump(); +} diff --git a/tslang/test/tester/jit-cache/sub/derived_module.ts b/tslang/test/tester/jit-cache/sub/derived_module.ts new file mode 100644 index 000000000..28255df47 --- /dev/null +++ b/tslang/test/tester/jit-cache/sub/derived_module.ts @@ -0,0 +1,14 @@ +import '../base_module' + +export class Derived extends Base { + hello() { + return "derived:" + super.hello(); + } +} + +export function bumpFromDerived() { + return bump(); +} + +// base_module's initializer has run already +export let doubled = counter * 2; diff --git a/tslang/tslang/CMakeLists.txt b/tslang/tslang/CMakeLists.txt index d1e6b941d..4cb4420cf 100644 --- a/tslang/tslang/CMakeLists.txt +++ b/tslang/tslang/CMakeLists.txt @@ -58,6 +58,7 @@ add_llvm_executable(tslang transform.cpp dump.cpp jit.cpp + jitcache.cpp obj.cpp exe.cpp TextDiagnostic.cpp diff --git a/tslang/tslang/jit.cpp b/tslang/tslang/jit.cpp index 0d1db3d0b..4e1b5c0cf 100644 --- a/tslang/tslang/jit.cpp +++ b/tslang/tslang/jit.cpp @@ -21,6 +21,7 @@ #include #include +#include #ifdef _WIN32 #include #include @@ -38,6 +39,8 @@ #include "TypeScript/Defines.h" +#include "jitcache.h" + // From TypeScript/ObjDumper.h, which cannot be included here: its llvm/BinaryFormat/COFF.h // collides with the IMAGE_* macros defines above. namespace Dump @@ -118,9 +121,15 @@ int loadLibrary(mlir::SmallString<256> &libPath, llvm::StringMap &export } #ifdef _WIN64 -// Image base of the JIT'd module (tslang JITs a single module per run), needed by -// the _CxxThrowException shim below. Filled in by JitSectionMemoryManager. -static uint64_t jitImageBase = 0; +// The JIT'd objects' address ranges and image bases, needed by the _CxxThrowException shim +// below: one for the program, and under the JIT cache one more for each module it imports, each +// with the base RTDyld resolved that object's RVAs against. Filled in by JitSectionMemoryManager. +struct JitImage +{ + uint64_t base; + uint64_t end; +}; +static std::deque jitImages; #endif // A SectionMemoryManager that makes JIT'd code behave like AOT'd code: @@ -177,7 +186,7 @@ class JitSectionMemoryManager : public llvm::SectionMemoryManager llvm::StringRef sectionName) override { auto *addr = llvm::SectionMemoryManager::allocateCodeSection(size, alignment, sectionID, sectionName); - noteSectionAddress(addr); + noteSectionAddress(addr, size); return addr; } @@ -193,7 +202,7 @@ class JitSectionMemoryManager : public llvm::SectionMemoryManager #endif auto *addr = llvm::SectionMemoryManager::allocateDataSection(size, alignment, sectionID, sectionName, isReadOnly); - noteSectionAddress(addr); + noteSectionAddress(addr, size); if (!isReadOnly && addr != nullptr) { if (auto addRoots = reinterpret_cast( @@ -246,7 +255,7 @@ class JitSectionMemoryManager : public llvm::SectionMemoryManager } private: - void noteSectionAddress(uint8_t *addr) + void noteSectionAddress(uint8_t *addr, uintptr_t size) { if (addr == nullptr) { @@ -256,13 +265,25 @@ class JitSectionMemoryManager : public llvm::SectionMemoryManager if (imageBase == 0 || reinterpret_cast(addr) < imageBase) { imageBase = reinterpret_cast(addr); + } + #ifdef _WIN64 - jitImageBase = imageBase; -#endif + // a memory manager serves one object: its image is the span of its sections + if (image == nullptr) + { + jitImages.push_back({imageBase, imageBase}); + image = &jitImages.back(); } + + image->base = imageBase; + image->end = std::max(image->end, reinterpret_cast(addr) + size); +#endif } uint64_t imageBase = 0; +#ifdef _WIN64 + JitImage *image = nullptr; +#endif mlir::SmallVector> gcRootSections; #ifdef _WIN64 mlir::SmallVector functionTables; @@ -356,10 +377,27 @@ static void jitCxxThrowException(void *exceptionObject, void *throwInfo) RtlPcToFileHeader(throwInfo, &moduleBase); } + // otherwise the JIT'd object the ThrowInfo is in + if (moduleBase == nullptr) + { + auto address = reinterpret_cast(throwInfo); + for (auto &image : jitImages) + { + if (address >= image.base && address < image.end) + { + moduleBase = reinterpret_cast(image.base); + break; + } + } + } + + if (moduleBase == nullptr && !jitImages.empty()) + { + moduleBase = reinterpret_cast(jitImages.front().base); + } + ULONG_PTR args[] = {cxxMagicNumber, reinterpret_cast(exceptionObject), - reinterpret_cast(throwInfo), - moduleBase != nullptr ? reinterpret_cast(moduleBase) - : static_cast(jitImageBase)}; + reinterpret_cast(throwInfo), reinterpret_cast(moduleBase)}; RaiseException(cxxExceptionCode, EXCEPTION_NONCONTINUABLE, 4, args); } @@ -403,7 +441,7 @@ static bool callEntryThunk(JitEntryThunkFn entryThunk, int argc, char **argv, in // converted to the first parameter's type, `argv` and `argc` as the `string[]` (data, length) of the // second or `argv` alone for a `Ref`, and the result converted to the exit code. LowerToAffineLoops has already refused any // other shape of `main`, with a source location; this still checks, for an entry picked with `-e`. -static llvm::Error addEntryThunk(llvm::Module &llvmModule, llvm::StringRef entryName) +llvm::Error addEntryThunk(llvm::Module &llvmModule, llvm::StringRef entryName) { auto *entry = llvmModule.getFunction(entryName); if (!entry) @@ -492,11 +530,12 @@ static llvm::Error addEntryThunk(llvm::Module &llvmModule, llvm::StringRef entry return llvm::Error::success(); } -int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compileOptions) + +// tslang.exe is an x64 process: it cannot execute i386 or wasm code in-process. Refused up front +// rather than left to LLJIT, whose failure for this is a relocation or "symbol not found" error +// deep in the session that never mentions the triple as the cause. +static bool checkInProcessJit(CompileOptions &compileOptions) { - // tslang.exe is an x64 process: it cannot execute i386 or wasm code in-process. Refused up - // front rather than left to LLJIT, whose failure for this is a relocation or "symbol not - // found" error deep in the session that never mentions the triple as the cause. if (!compileOptions.targetInfo.supportsInProcessJit) { llvm::WithColor::error(llvm::errs(), "tslang") @@ -504,28 +543,16 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile << llvm::sys::getDefaultTargetTriple() << ", so it cannot run code built for " << compileOptions.moduleTargetTriple << ". Build it instead with --emit=exe or --emit=obj.\n"; - return -1; + return false; } - // to avoid false positive memory leak reports in release builds - // Print a stack trace if we signal out. - llvm::sys::PrintStackTraceOnErrorSignal(argv[0]); - - llvm::PrettyStackTraceProgram X(argc, argv); - llvm::setBugReportMsg("PLEASE submit a bug report to https://github.com/ASDAlexander77/TypeScriptCompiler/issues and include the crash backtrace."); - - llvm::llvm_shutdown_obj Y; // Call llvm_shutdown() on exit. - - // No leaks until here - - registerMLIRDialects(module); - - // Initialize LLVM targets. - llvm::InitializeNativeTarget(); - llvm::InitializeNativeTargetAsmPrinter(); - - auto optPipeline = getTransformer(enableOpt, optLevel, sizeLevel, compileOptions); + return true; +} +// The libraries the program runs against - the default library, the GC runtime - added to +// --shared-libs, and on Windows the CRT entry points bound to this process's own CRT. +static int prepareJitProcess(CompileOptions &compileOptions) +{ // If shared library implements custom mlir-runner library init and destroy // functions, we'll use them to register the library with the execution // engine. Otherwise we'll pass library directly to the execution engine. @@ -668,64 +695,13 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile } #endif - if (dumpObjectFile) - { - // Compile-and-dump only, no execution: the stock MLIR engine is enough. - mlir::SmallVector sharedLibPaths; - sharedLibPaths.append(begin(clSharedLibs), end(clSharedLibs)); - - mlir::ExecutionEngineOptions engineOptions; - engineOptions.transformer = optPipeline; - engineOptions.enableObjectDump = true; - engineOptions.enableGDBNotificationListener = !enableOpt; - engineOptions.sharedLibPaths = sharedLibPaths; - if (enableOpt.getValue()) - { - engineOptions.jitCodeGenOptLevel = (llvm::CodeGenOptLevel) optLevel.getValue(); - } - - auto maybeEngine = mlir::ExecutionEngine::create(module, engineOptions); - if (!maybeEngine) - { - auto err = maybeEngine.takeError(); - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to construct an execution engine, error: " << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; - } - auto &engine = maybeEngine.get(); - - auto expectedFPtr = engine->lookup(mainFuncName); - if (!expectedFPtr) - { - auto err = expectedFPtr.takeError(); - llvm::WithColor::error(llvm::errs(), "tslang") << err; - llvm::consumeError(std::move(err)); - return -1; - } - - llvm::Triple TheTriple; - std::string targetTriple = llvm::sys::getDefaultTargetTriple(); - if (!TargetTriple.empty()) - { - targetTriple = llvm::Triple::normalize(TargetTriple); - } - - TheTriple = llvm::Triple(targetTriple); - - engine->dumpToObjectFile( - objectFilename.empty() - ? inputFilename + ((TheTriple.getOS() == llvm::Triple::Win32) ? ".obj" : ".o") - : objectFilename); - - return 0; - } - - // Run path: build our own LLJIT instead of mlir::ExecutionEngine — the stock - // engine hard-codes a plain SectionMemoryManager, which neither registers GC - // roots for JIT'd globals nor Win64 unwind info (see JitSectionMemoryManager). + return 0; +} - // Load the shared libraries into the process; the JIT resolves external - // symbols from their export tables via the current-process generator below. +// Load the shared libraries into the process; the JIT resolves external symbols from their +// export tables via the current-process generator (see createJit). +static int loadJitSharedLibraries(bool &needsGCEnableThreadsStandIn) +{ for (auto &libPathStr : clSharedLibs) { std::string errMsg; @@ -748,24 +724,20 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile // the answer decides whether to shadow a real definition: SearchForAddressOfSymbol sees the // export tables of everything loaded so far, which is exactly what the JIT's process generator // will see. - auto needsGCEnableThreadsStandIn = + needsGCEnableThreadsStandIn = llvm::sys::DynamicLibrary::SearchForAddressOfSymbol("GC_enable_threads") == nullptr; - auto llvmContext = std::make_unique(); - auto llvmModule = mlir::translateModuleToLLVMIR(module, *llvmContext); - if (!llvmModule) - { - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to emit LLVM IR\n"; - return -1; - } + return 0; +} +// The host's target machine builder, at the optimization level asked for; the JIT and the JIT +// cache compile with it, so a cached object is what the JIT itself would have made. +llvm::Expected getJitTargetMachineBuilder() +{ auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost(); if (!tmBuilderOrError) { - auto err = tmBuilderOrError.takeError(); - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a JITTargetMachineBuilder for the host, error: " << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; + return tmBuilderOrError.takeError(); } if (enableOpt.getValue()) @@ -773,39 +745,20 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile tmBuilderOrError->setCodeGenOptLevel((llvm::CodeGenOptLevel) optLevel.getValue()); } - auto tmOrError = tmBuilderOrError->createTargetMachine(); - if (!tmOrError) - { - auto err = tmOrError.takeError(); - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a TargetMachine for the host, error: " << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; - } - - llvmModule->setDataLayout((*tmOrError)->createDataLayout()); - llvmModule->setTargetTriple((*tmOrError)->getTargetTriple()); - - if (auto err = optPipeline(llvmModule.get())) - { - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to optimize LLVM IR, error: " << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; - } - - // after the optimizer, so it cannot inline an entry with debug info into a thunk without any - if (auto err = addEntryThunk(*llvmModule, mainFuncName.getValue())) - { - llvm::WithColor::error(llvm::errs(), "tslang") << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; - } + return tmBuilderOrError; +} +// Our own LLJIT instead of mlir::ExecutionEngine — the stock engine hard-codes a plain +// SectionMemoryManager, which neither registers GC roots for JIT'd globals nor Win64 unwind info +// (see JitSectionMemoryManager). Symbols the program does not define come from this process. +static std::unique_ptr createJit(llvm::orc::JITTargetMachineBuilder tmBuilder, bool needsGCEnableThreadsStandIn) +{ + auto targetTriple = tmBuilder.getTargetTriple(); auto maybeJit = llvm::orc::LLJITBuilder() - .setJITTargetMachineBuilder(std::move(*tmBuilderOrError)) - .setDataLayout(llvmModule->getDataLayout()) + .setJITTargetMachineBuilder(std::move(tmBuilder)) .setObjectLinkingLayerCreator( - [targetTriple = llvmModule->getTargetTriple()](llvm::orc::ExecutionSession &session) + [targetTriple](llvm::orc::ExecutionSession &session) -> llvm::Expected> { auto objectLayer = std::make_unique( session, [](const llvm::MemoryBuffer &) { return std::make_unique(); }); @@ -834,10 +787,10 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile auto err = maybeJit.takeError(); llvm::WithColor::error(llvm::errs(), "tslang") << "failed to construct the JIT engine, error: " << err << "\n"; llvm::consumeError(std::move(err)); - return -1; + return nullptr; } - auto &jit = maybeJit.get(); + auto jit = std::move(maybeJit.get()); // Resolve symbols from the current process, including the loaded shared // libraries and the AddSymbol overrides above. @@ -847,7 +800,7 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile auto err = generator.takeError(); llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a process symbol generator, error: " << err << "\n"; llvm::consumeError(std::move(err)); - return -1; + return nullptr; } jit->getMainJITDylib().addGenerator(std::move(*generator)); @@ -864,7 +817,7 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile { llvm::WithColor::error(llvm::errs(), "tslang") << "failed to define the shared library loader, error: " << err << "\n"; llvm::consumeError(std::move(err)); - return -1; + return nullptr; } } @@ -911,7 +864,7 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile { llvm::WithColor::error(llvm::errs(), "tslang") << "failed to define CRT overrides, error: " << err << "\n"; llvm::consumeError(std::move(err)); - return -1; + return nullptr; } } #else @@ -925,20 +878,21 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile { llvm::WithColor::error(llvm::errs(), "tslang") << "failed to define the assert handler, error: " << err << "\n"; llvm::consumeError(std::move(err)); - return -1; + return nullptr; } } #endif - if (auto err = jit->addIRModule(llvm::orc::ThreadSafeModule(std::move(llvmModule), std::move(llvmContext)))) - { - llvm::WithColor::error(llvm::errs(), "tslang") << "failed to add the module to the JIT engine, error: " << err << "\n"; - llvm::consumeError(std::move(err)); - return -1; - } + return jit; +} +// Runs what was added to the JIT: the platform initializers, the objects' own initializers +// (`initFunctions`, in the order given - an imported module's before its importer's), the +// program's `__mlir_gctors` if it has one, then the entry point. +static int runJitProgram(llvm::orc::LLJIT &jit, llvm::ArrayRef initFunctions, bool hasGCtorsMethod) +{ // run platform initializers (llvm.global_ctors etc.) - if (auto err = jit->initialize(jit->getMainJITDylib())) + if (auto err = jit.initialize(jit.getMainJITDylib())) { llvm::WithColor::error(llvm::errs(), "tslang") << "JIT initialization failed, error: " << err << "\n"; llvm::consumeError(std::move(err)); @@ -946,7 +900,7 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile } auto invoke = [&](llvm::StringRef name) { - auto sym = jit->lookup(name); + auto sym = jit.lookup(name); if (!sym) { // Streaming an Error only logs it - the payload survives, and ~Error then trips @@ -963,12 +917,20 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile return 0; }; - if (module.lookupSymbol(MLIR_GCTORS) && invoke(MLIR_GCTORS) != 0) + for (auto &initFunction : initFunctions) + { + if (invoke(initFunction) != 0) + { + return -1; + } + } + + if (hasGCtorsMethod && invoke(MLIR_GCTORS) != 0) { return -1; } - auto entryThunk = jit->lookup(JIT_ENTRY_THUNK_NAME); + auto entryThunk = jit.lookup(JIT_ENTRY_THUNK_NAME); if (!entryThunk) { auto err = entryThunk.takeError(); @@ -1029,3 +991,226 @@ int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compile #endif return exitCode; } + +int runJit(int argc, char **argv, mlir::ModuleOp module, CompileOptions &compileOptions) +{ + if (!checkInProcessJit(compileOptions)) + { + return -1; + } + + // to avoid false positive memory leak reports in release builds + // Print a stack trace if we signal out. + llvm::sys::PrintStackTraceOnErrorSignal(argv[0]); + + llvm::PrettyStackTraceProgram X(argc, argv); + llvm::setBugReportMsg("PLEASE submit a bug report to https://github.com/ASDAlexander77/TypeScriptCompiler/issues and include the crash backtrace."); + + llvm::llvm_shutdown_obj Y; // Call llvm_shutdown() on exit. + + // No leaks until here + + registerMLIRDialects(module); + + // Initialize LLVM targets. + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + + auto optPipeline = getTransformer(enableOpt, optLevel, sizeLevel, compileOptions); + + if (auto result = prepareJitProcess(compileOptions)) + { + return result; + } + + if (dumpObjectFile) + { + // Compile-and-dump only, no execution: the stock MLIR engine is enough. + mlir::SmallVector sharedLibPaths; + sharedLibPaths.append(begin(clSharedLibs), end(clSharedLibs)); + + mlir::ExecutionEngineOptions engineOptions; + engineOptions.transformer = optPipeline; + engineOptions.enableObjectDump = true; + engineOptions.enableGDBNotificationListener = !enableOpt; + engineOptions.sharedLibPaths = sharedLibPaths; + if (enableOpt.getValue()) + { + engineOptions.jitCodeGenOptLevel = (llvm::CodeGenOptLevel) optLevel.getValue(); + } + + auto maybeEngine = mlir::ExecutionEngine::create(module, engineOptions); + if (!maybeEngine) + { + auto err = maybeEngine.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to construct an execution engine, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + auto &engine = maybeEngine.get(); + + auto expectedFPtr = engine->lookup(mainFuncName); + if (!expectedFPtr) + { + auto err = expectedFPtr.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << err; + llvm::consumeError(std::move(err)); + return -1; + } + + llvm::Triple TheTriple; + std::string targetTriple = llvm::sys::getDefaultTargetTriple(); + if (!TargetTriple.empty()) + { + targetTriple = llvm::Triple::normalize(TargetTriple); + } + + TheTriple = llvm::Triple(targetTriple); + + engine->dumpToObjectFile( + objectFilename.empty() + ? inputFilename + ((TheTriple.getOS() == llvm::Triple::Win32) ? ".obj" : ".o") + : objectFilename); + + return 0; + } + + auto needsGCEnableThreadsStandIn = false; + if (auto result = loadJitSharedLibraries(needsGCEnableThreadsStandIn)) + { + return result; + } + + auto llvmContext = std::make_unique(); + auto llvmModule = mlir::translateModuleToLLVMIR(module, *llvmContext); + if (!llvmModule) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to emit LLVM IR\n"; + return -1; + } + + auto tmBuilderOrError = getJitTargetMachineBuilder(); + if (!tmBuilderOrError) + { + auto err = tmBuilderOrError.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a JITTargetMachineBuilder for the host, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + auto tmOrError = tmBuilderOrError->createTargetMachine(); + if (!tmOrError) + { + auto err = tmOrError.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a TargetMachine for the host, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + llvmModule->setDataLayout((*tmOrError)->createDataLayout()); + llvmModule->setTargetTriple((*tmOrError)->getTargetTriple()); + + if (auto err = optPipeline(llvmModule.get())) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to optimize LLVM IR, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + // after the optimizer, so it cannot inline an entry with debug info into a thunk without any + if (auto err = addEntryThunk(*llvmModule, mainFuncName.getValue())) + { + llvm::WithColor::error(llvm::errs(), "tslang") << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + auto jit = createJit(std::move(*tmBuilderOrError), needsGCEnableThreadsStandIn); + if (!jit) + { + return -1; + } + + if (auto err = jit->addIRModule(llvm::orc::ThreadSafeModule(std::move(llvmModule), std::move(llvmContext)))) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to add the module to the JIT engine, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + return runJitProgram(*jit, {}, module.lookupSymbol(MLIR_GCTORS) != nullptr); +} + +// The JIT with its cache: the program and every .ts module it imports are each an object of +// their own, compiled - or found already compiled - by buildJitUnits (jitcache.cpp), and loaded +// side by side, as the objects of a compiled program are linked. +int runJitCached(int argc, char **argv, CompileOptions &compileOptions) +{ + if (!checkInProcessJit(compileOptions)) + { + return -1; + } + + llvm::sys::PrintStackTraceOnErrorSignal(argv[0]); + + llvm::PrettyStackTraceProgram X(argc, argv); + llvm::setBugReportMsg("PLEASE submit a bug report to https://github.com/ASDAlexander77/TypeScriptCompiler/issues and include the crash backtrace."); + + llvm::llvm_shutdown_obj Y; // Call llvm_shutdown() on exit. + + llvm::InitializeNativeTarget(); + llvm::InitializeNativeTargetAsmPrinter(); + + if (auto result = prepareJitProcess(compileOptions)) + { + return result; + } + + auto needsGCEnableThreadsStandIn = false; + if (auto result = loadJitSharedLibraries(needsGCEnableThreadsStandIn)) + { + return result; + } + + auto tmBuilderOrError = getJitTargetMachineBuilder(); + if (!tmBuilderOrError) + { + auto err = tmBuilderOrError.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a JITTargetMachineBuilder for the host, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + std::vector units; + if (auto result = buildJitUnits(argv[0], inputFilename.getValue(), compileOptions, *tmBuilderOrError, units)) + { + return result; + } + + auto jit = createJit(std::move(*tmBuilderOrError), needsGCEnableThreadsStandIn); + if (!jit) + { + return -1; + } + + std::vector initFunctions; + auto hasGCtorsMethod = false; + for (auto &unit : units) + { + if (auto err = jit->addObjectFile(std::move(unit.object))) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to add '" << unit.sourcePath << "' to the JIT engine, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + if (!unit.initFunction.empty()) + { + initFunctions.push_back(unit.initFunction); + } + + hasGCtorsMethod |= unit.isMain && unit.hasGCtorsMethod; + } + + return runJitProgram(*jit, initFunctions, hasGCtorsMethod); +} diff --git a/tslang/tslang/jitcache.cpp b/tslang/tslang/jitcache.cpp new file mode 100644 index 000000000..7672659f0 --- /dev/null +++ b/tslang/tslang/jitcache.cpp @@ -0,0 +1,637 @@ +// The JIT cache: every .ts file a JIT run needs - the program and each .ts module it imports - is +// compiled into an object of its own and kept in a cache folder, `__jit` beside the source (or +// the folder given with --jit-cache-dir). The next run loads the object instead of compiling the +// file again, as long as nothing it was compiled from has changed. +// +// Next to each object is a manifest (.deps) saying what the object was compiled from: +// +// tslang-jit-cache 1 +// compiler +// object +// init +// gctors <1 if the file has __mlir_gctors> +// import +// dep +// +// An object is used only if every one of these still holds. The name of an object is the file's +// name and a hash of its path and the options it was compiled with, so a run with other options +// has objects of its own and does not replace these. + +#include "TypeScript/DataStructs.h" +#include "TypeScript/Defines.h" + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OwningOpRef.h" +#include "mlir/Target/LLVMIR/Export.h" + +#include "llvm/ADT/SmallString.h" +#include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/StringMap.h" +#include "llvm/ADT/StringSet.h" +#include "llvm/ExecutionEngine/Orc/CompileUtils.h" +#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/Program.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/WithColor.h" +#include "llvm/Support/xxhash.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Target/TargetMachine.h" + +#include "jitcache.h" + +#include + +#define DEBUG_TYPE "tslang" + +namespace cl = llvm::cl; + +extern cl::opt enableOpt; +extern cl::opt optLevel; +extern cl::opt sizeLevel; +extern cl::opt mainFuncName; +extern cl::opt verifyOwnership; +extern cl::opt ownSkipInference; +extern cl::opt jitCacheDir; + +std::unique_ptr createMLIRContext(); +int compileTypeScriptFileIntoMLIR(mlir::MLIRContext &, llvm::StringRef, llvm::SourceMgr &, mlir::OwningOpRef &, CompileOptions &); +int runMLIRPasses(mlir::MLIRContext &, llvm::SourceMgr &, mlir::OwningOpRef &, CompileOptions &); +int registerMLIRDialects(mlir::ModuleOp); +std::function getTransformer(bool, int, int, CompileOptions &); +llvm::Error addEntryThunk(llvm::Module &, llvm::StringRef); + +#define JIT_CACHE_FORMAT "tslang-jit-cache 1" +#define JIT_CACHE_DIR "__jit" +#define JIT_INIT_PREFIX "__tslang_jit_init_" + +namespace +{ + +std::string hashString(llvm::StringRef data) +{ + return llvm::utohexstr(llvm::xxh3_64bits(data), /*LowerCase=*/true, /*Width=*/16); +} + +std::string canonicalPath(llvm::StringRef path) +{ + llvm::SmallString<256> canonical; + if (!llvm::sys::fs::real_path(path, canonical)) + { + return canonical.str().str(); + } + + canonical = path; + llvm::sys::fs::make_absolute(canonical); + llvm::sys::path::remove_dots(canonical, /*remove_dot_dot=*/true); + return canonical.str().str(); +} + +// Written beside the final name and renamed over it, so a run reading the cache - another test +// running the same module at the same time, say - never sees half a file. +bool writeFileAtomically(llvm::StringRef path, llvm::StringRef content) +{ + int fd; + llvm::SmallString<256> tempPath; + if (llvm::sys::fs::createUniqueFile(path + ".%%%%%%%%.tmp", fd, tempPath)) + { + return false; + } + + { + llvm::raw_fd_ostream os(fd, /*shouldClose=*/true); + os << content; + os.close(); + if (os.has_error()) + { + os.clear_error(); + llvm::sys::fs::remove(tempPath); + return false; + } + } + + if (llvm::sys::fs::rename(tempPath, path)) + { + llvm::sys::fs::remove(tempPath); + return false; + } + + return true; +} + +// The JIT runs no object's llvm.global_ctors - only an IR module's, whose constructors its +// platform takes out before compiling it. So they are taken out here, into one function, in the +// order of their priority; the JIT calls it before the program starts (runJitProgram). +std::string moveGlobalCtorsIntoFunction(llvm::Module &llvmModule, llvm::StringRef functionName) +{ + auto *globalCtors = llvmModule.getNamedGlobal("llvm.global_ctors"); + if (!globalCtors) + { + return ""; + } + + struct Ctor + { + unsigned priority; + llvm::Function *function; + }; + + std::vector ctors; + for (auto ctor : llvm::orc::getConstructors(llvmModule)) + { + if (ctor.Func) + { + ctors.push_back({ctor.Priority, ctor.Func}); + } + } + + globalCtors->eraseFromParent(); + if (ctors.empty()) + { + return ""; + } + + std::stable_sort(ctors.begin(), ctors.end(), [](const Ctor &a, const Ctor &b) { return a.priority < b.priority; }); + + auto &context = llvmModule.getContext(); + auto *initFunction = llvm::Function::Create(llvm::FunctionType::get(llvm::Type::getVoidTy(context), false), + llvm::Function::ExternalLinkage, functionName, llvmModule); + // a constructor may throw through it; Win64 can only unwind a frame it has unwind info for + initFunction->setUWTableKind(llvm::UWTableKind::Async); + + llvm::IRBuilder<> builder(llvm::BasicBlock::Create(context, "entry", initFunction)); + for (auto &ctor : ctors) + { + auto *call = builder.CreateCall(ctor.function->getFunctionType(), ctor.function); + call->setCallingConv(ctor.function->getCallingConv()); + } + + builder.CreateRetVoid(); + return functionName.str(); +} + +// A definition several objects have - a class's `.size`, say, in the module declaring the class +// and in every module importing it - is put in a COMDAT, and a linker keeps one of them. RTDyld +// does that for COFF (ORC takes a COMDAT symbol for a weak one) but not for ELF, where the second +// object loaded fails with "duplicate definition". Made weak, ORC keeps the first. +void weakenComdatDefinitions(llvm::Module &llvmModule) +{ + if (!llvmModule.getTargetTriple().isOSBinFormatELF()) + { + return; + } + + for (auto &globalObject : llvmModule.global_objects()) + { + if (globalObject.hasComdat() && !globalObject.isDeclaration() && globalObject.hasExternalLinkage()) + { + globalObject.setLinkage(llvm::GlobalValue::WeakAnyLinkage); + } + } +} + +struct Manifest +{ + std::string compiler; + std::string objectHash; + std::string initFunction; + bool hasGCtorsMethod = false; + std::vector imports; + // hash, path + std::vector> deps; + + std::string str() const + { + std::string text; + llvm::raw_string_ostream os(text); + os << JIT_CACHE_FORMAT "\n"; + os << "compiler " << compiler << "\n"; + os << "object " << objectHash << "\n"; + os << "init " << (initFunction.empty() ? "-" : initFunction) << "\n"; + os << "gctors " << (hasGCtorsMethod ? 1 : 0) << "\n"; + for (auto &import : imports) + { + os << "import " << import << "\n"; + } + + for (auto &[hash, path] : deps) + { + os << "dep " << hash << " " << path << "\n"; + } + + return text; + } + + bool parse(llvm::StringRef text) + { + llvm::SmallVector lines; + text.split(lines, '\n', /*MaxSplit=*/-1, /*KeepEmpty=*/false); + if (lines.empty() || lines.front().rtrim() != JIT_CACHE_FORMAT) + { + return false; + } + + for (auto line : llvm::ArrayRef(lines).drop_front()) + { + line = line.rtrim("\r"); + auto [key, value] = line.split(' '); + if (key == "compiler") + { + compiler = value.str(); + } + else if (key == "object") + { + objectHash = value.str(); + } + else if (key == "init") + { + initFunction = value == "-" ? "" : value.str(); + } + else if (key == "gctors") + { + hasGCtorsMethod = value == "1"; + } + else if (key == "import") + { + imports.push_back(value.str()); + } + else if (key == "dep") + { + auto [hash, path] = value.split(' '); + deps.push_back({hash.str(), path.str()}); + } + else + { + return false; + } + } + + return !compiler.empty() && !objectHash.empty(); + } +}; + +class JitUnitBuilder +{ + public: + JitUnitBuilder(const char *argv0, const CompileOptions &compileOptions, llvm::orc::JITTargetMachineBuilder &tmBuilder, + std::vector &units) + : compileOptions(compileOptions), tmBuilder(tmBuilder), units(units) + { + // a rebuilt compiler compiles differently: its objects are not these + auto executable = llvm::sys::fs::getMainExecutable(argv0, reinterpret_cast(&canonicalPath)); + llvm::sys::fs::file_status status; + if (!executable.empty() && !llvm::sys::fs::status(executable, status)) + { + compilerId = llvm::utohexstr(status.getSize(), true) + "-" + + llvm::utohexstr(status.getLastModificationTime().time_since_epoch().count(), true); + } + } + + int add(llvm::StringRef sourcePath, bool isMain) + { + auto path = canonicalPath(sourcePath); + if (!visited.insert(path).second) + { + // loaded already, or on the way: an import cycle, or a module importing the program + return 0; + } + + if (!tm) + { + auto tmOrError = tmBuilder.createTargetMachine(); + if (!tmOrError) + { + auto err = tmOrError.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to create a TargetMachine for the host, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + tm = std::move(*tmOrError); + } + + JitUnit unit; + unit.sourcePath = path; + unit.isMain = isMain; + + auto [objectPath, manifestPath] = cachePaths(path, isMain); + + Manifest manifest; + if (!loadFromCache(objectPath, manifestPath, manifest, unit)) + { + if (auto result = compile(path, isMain, manifest, unit)) + { + return result; + } + + storeInCache(objectPath, manifestPath, manifest, *unit.object); + } + + unit.initFunction = manifest.initFunction; + unit.hasGCtorsMethod = manifest.hasGCtorsMethod; + + // what it imports first: their objects are loaded, and their constructors run, before its own + for (auto &import : manifest.imports) + { + if (auto result = add(import, /*isMain=*/false)) + { + return result; + } + } + + units.push_back(std::move(unit)); + return 0; + } + + private: + // Everything that changes what a file compiles into, but the files themselves (the manifest's + // deps) and the compiler (its `compiler`). + std::string optionsKey(bool isMain) + { + std::string key; + llvm::raw_string_ostream os(key); + os << (isMain ? "main " + mainFuncName.getValue() : std::string("import")) << "\n"; + os << "mm=" << compileOptions.memoryModel << " opt=" << enableOpt.getValue() << " opt_level=" << optLevel.getValue() + << " size_level=" << sizeLevel.getValue() << " di=" << compileOptions.generateDebugInfo + << " lldb=" << compileOptions.lldbDebugInfo << " no-default-lib=" << compileOptions.noDefaultLib + << " builtins=" << compileOptions.enableBuiltins << " strict-null-checks=" << compileOptions.strictNullChecks + << " fast-math=" << compileOptions.enableFastMath << " gctors-as-method=" << compileOptions.appendGCtorsToMethod + << " embed-declarations=" << compileOptions.embedExportDeclarations + << " verify-ownership=" << verifyOwnership.getValue() << " own-skip-inference=" << ownSkipInference.getValue() + << "\n"; + os << "export=" << llvm::join(compileOptions.exportFilters, ",") << "\n"; + os << "default-lib=" << compileOptions.defaultDeclarationTSFile << "\n"; + os << "target=" << tmBuilder.getTargetTriple().str() << " cpu=" << tmBuilder.getCPU() + << " features=" << tmBuilder.getFeatures().getString() << "\n"; + return key; + } + + std::pair cachePaths(llvm::StringRef path, bool isMain) + { + llvm::SmallString<256> objectPath; + if (jitCacheDir.empty()) + { + objectPath = llvm::sys::path::parent_path(path); + llvm::sys::path::append(objectPath, JIT_CACHE_DIR); + } + else + { + objectPath = jitCacheDir.getValue(); + } + + auto name = llvm::sys::path::filename(path).str() + "." + hashString(path.str() + "\n" + optionsKey(isMain)); + name += tm->getTargetTriple().isOSBinFormatCOFF() ? ".obj" : ".o"; + llvm::sys::path::append(objectPath, name); + + auto objectPathStr = objectPath.str().str(); + return {objectPathStr, objectPathStr + ".deps"}; + } + + // the hash of a file's content, or empty if it cannot be read; once per file and run + std::string fileHash(llvm::StringRef path) + { + auto it = fileHashes.find(path); + if (it != fileHashes.end()) + { + return it->second; + } + + std::string hash; + auto fileOrErr = llvm::MemoryBuffer::getFile(path, /*IsText=*/false, /*RequiresNullTerminator=*/false); + if (fileOrErr) + { + hash = hashString((*fileOrErr)->getBuffer()); + } + + fileHashes[path] = hash; + return hash; + } + + bool loadFromCache(llvm::StringRef objectPath, llvm::StringRef manifestPath, Manifest &manifest, JitUnit &unit) + { + auto manifestOrErr = llvm::MemoryBuffer::getFile(manifestPath, /*IsText=*/true); + if (!manifestOrErr || !manifest.parse((*manifestOrErr)->getBuffer())) + { + manifest = Manifest(); + return false; + } + + auto stale = [&](llvm::StringRef why) { + LLVM_DEBUG(llvm::dbgs() << "JIT cache: " << objectPath << " is out of date: " << why << "\n";); + (void)why; + manifest = Manifest(); + return false; + }; + + if (compilerId.empty() || manifest.compiler != compilerId) + { + return stale("compiled by another tslang"); + } + + for (auto &[hash, path] : manifest.deps) + { + if (fileHash(path) != hash) + { + return stale(path); + } + } + + // Volatile: read, not mapped, so another run can still replace the file while this one runs + auto objectOrErr = llvm::MemoryBuffer::getFile(objectPath, /*IsText=*/false, /*RequiresNullTerminator=*/false, + /*IsVolatile=*/true); + if (!objectOrErr || hashString((*objectOrErr)->getBuffer()) != manifest.objectHash) + { + return stale("the object is missing or is not the one compiled"); + } + + LLVM_DEBUG(llvm::dbgs() << "JIT cache: using " << objectPath << "\n";); + unit.object = std::move(*objectOrErr); + return true; + } + + void storeInCache(llvm::StringRef objectPath, llvm::StringRef manifestPath, const Manifest &manifest, + const llvm::MemoryBuffer &object) + { + // Without a compiler to tell one build from another, the object could outlive the build it + // belongs to: better compiled again every time. + if (compilerId.empty()) + { + return; + } + + // Nowhere to write it (a read-only folder, say) is not an error: the program still runs, + // it is compiled again next time. + if (llvm::sys::fs::create_directories(llvm::sys::path::parent_path(objectPath)) || + !writeFileAtomically(objectPath, object.getBuffer()) || !writeFileAtomically(manifestPath, manifest.str())) + { + LLVM_DEBUG(llvm::dbgs() << "JIT cache: could not write " << objectPath << "\n";); + return; + } + + LLVM_DEBUG(llvm::dbgs() << "JIT cache: wrote " << objectPath << "\n";); + } + + int compile(llvm::StringRef path, bool isMain, Manifest &manifest, JitUnit &unit) + { + LLVM_DEBUG(llvm::dbgs() << "JIT cache: compiling " << path << "\n";); + + // An imported module is compiled as `--emit=obj` compiles a module linked beside the + // program: no entry point, its top level run from its global constructors, and what it + // imports declared only - those have objects of their own. + auto options = compileOptions; + if (!isMain) + { + options.isJit = false; + options.generateEntryPoint = false; + } + + std::unique_ptr context; + std::unique_ptr sourceMgr; + mlir::OwningOpRef module; + auto generate = [&]() { + options.sourceImports.clear(); + options.sharedLibraryImports.clear(); + options.importCycle = false; + module = nullptr; + context = createMLIRContext(); + sourceMgr = std::make_unique(); + return compileTypeScriptFileIntoMLIR(*context, path, *sourceMgr, module, options); + }; + + if (auto error = generate()) + { + return error; + } + + // A module in an import cycle cannot be compiled by itself when the other side of the + // cycle has code at its top level - the same holds for `--emit=obj`. So a program whose + // imports lead back to a file on the way is compiled as the JIT without the cache + // compiles it: one module, the imported files included with their bodies, an object that + // needs no other. + if (isMain && options.importCycle) + { + LLVM_DEBUG(llvm::dbgs() << "JIT cache: " << path << " has an import cycle, compiled as one module\n";); + options.jitCache = false; + if (auto error = generate()) + { + return error; + } + + options.sourceImports.clear(); + } + + if (auto error = runMLIRPasses(*context, *sourceMgr, module, options)) + { + return error; + } + + manifest.compiler = compilerId; + manifest.imports = options.sourceImports; + manifest.hasGCtorsMethod = isMain && (*module).lookupSymbol(MLIR_GCTORS) != nullptr; + + // What it was compiled from: every source file read - itself, the files it references, + // the modules it imports (their declarations are compiled into it), lib.d.ts - and the + // libraries it imports, whose declarations it was compiled against. + llvm::StringSet<> depPaths; + auto addDep = [&](llvm::StringRef depPath, std::string hash) { + if (!hash.empty() && depPaths.insert(depPath).second) + { + manifest.deps.push_back({hash, depPath.str()}); + } + }; + + for (unsigned id = 1; id <= sourceMgr->getNumBuffers(); id++) + { + auto *buffer = sourceMgr->getMemoryBuffer(id); + llvm::SmallString<256> depPath(buffer->getBufferIdentifier()); + llvm::sys::fs::make_absolute(depPath); + // only files: generated sources (a library's declarations, say) come from a file named below + if (!llvm::sys::fs::is_regular_file(depPath)) + { + continue; + } + + // what was read is what it was compiled from, even if the file changed since + auto hash = hashString(buffer->getBuffer()); + addDep(depPath, hash); + } + + for (auto &library : options.sharedLibraryImports) + { + addDep(library, fileHash(library)); + } + + registerMLIRDialects(*module); + + llvm::LLVMContext llvmContext; + auto llvmModule = mlir::translateModuleToLLVMIR(*module, llvmContext); + if (!llvmModule) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to emit LLVM IR\n"; + return -1; + } + + llvmModule->setDataLayout(tm->createDataLayout()); + llvmModule->setTargetTriple(tm->getTargetTriple()); + + auto optPipeline = getTransformer(enableOpt, optLevel, sizeLevel, options); + if (auto err = optPipeline(llvmModule.get())) + { + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to optimize LLVM IR, error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + if (isMain) + { + // after the optimizer, so it cannot inline an entry with debug info into a thunk without any + if (auto err = addEntryThunk(*llvmModule, mainFuncName.getValue())) + { + llvm::WithColor::error(llvm::errs(), "tslang") << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + } + + weakenComdatDefinitions(*llvmModule); + manifest.initFunction = moveGlobalCtorsIntoFunction(*llvmModule, JIT_INIT_PREFIX + hashString(path)); + + llvm::orc::SimpleCompiler compiler(*tm); + auto objectOrErr = compiler(*llvmModule); + if (!objectOrErr) + { + auto err = objectOrErr.takeError(); + llvm::WithColor::error(llvm::errs(), "tslang") << "failed to compile '" << path << "', error: " << err << "\n"; + llvm::consumeError(std::move(err)); + return -1; + } + + unit.object = std::move(*objectOrErr); + manifest.objectHash = hashString(unit.object->getBuffer()); + return 0; + } + + const CompileOptions &compileOptions; + llvm::orc::JITTargetMachineBuilder &tmBuilder; + std::unique_ptr tm; + std::vector &units; + std::string compilerId; + llvm::StringSet<> visited; + llvm::StringMap fileHashes; +}; + +} // namespace + +int buildJitUnits(const char *argv0, llvm::StringRef mainFile, const CompileOptions &compileOptions, + llvm::orc::JITTargetMachineBuilder &tmBuilder, std::vector &units) +{ + JitUnitBuilder builder(argv0, compileOptions, tmBuilder, units); + return builder.add(mainFile, /*isMain=*/true); +} diff --git a/tslang/tslang/jitcache.h b/tslang/tslang/jitcache.h new file mode 100644 index 000000000..f6e8e1080 --- /dev/null +++ b/tslang/tslang/jitcache.h @@ -0,0 +1,34 @@ +#ifndef TSLANG_JITCACHE_H_ +#define TSLANG_JITCACHE_H_ + +#include "TypeScript/DataStructs.h" + +#include "llvm/ADT/StringRef.h" +#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" +#include "llvm/Support/MemoryBuffer.h" + +#include +#include +#include + +// One .ts file compiled for the JIT: the program, or a module it imports. +struct JitUnit +{ + std::string sourcePath; + bool isMain = false; + std::unique_ptr object; + // what the file's global constructors became: the JIT runs no object's llvm.global_ctors, so + // they are called from one function, which the JIT calls before the program starts + std::string initFunction; + // `--gctors-as-method`: the program's __mlir_gctors + bool hasGCtorsMethod = false; +}; + +// The program in `mainFile` and every .ts module it imports, directly or through another, each an +// object: read from the cache if its sources and options have not changed since it was compiled, +// otherwise compiled and put in the cache. In `units` an imported module comes before its +// importer, the program last. Non-zero if a file does not compile. +int buildJitUnits(const char *argv0, llvm::StringRef mainFile, const CompileOptions &compileOptions, + llvm::orc::JITTargetMachineBuilder &tmBuilder, std::vector &units); + +#endif // TSLANG_JITCACHE_H_ diff --git a/tslang/tslang/opts.cpp b/tslang/tslang/opts.cpp index 0023d2601..cf80fb928 100644 --- a/tslang/tslang/opts.cpp +++ b/tslang/tslang/opts.cpp @@ -30,6 +30,8 @@ extern cl::opt entryPoint; extern cl::opt strictNullChecks; extern cl::opt embedExportDeclarationsAction; extern cl::opt enableFastMath; +extern cl::opt jitCache; +extern cl::opt dumpObjectFile; // obj extern cl::opt TargetTriple; @@ -66,6 +68,9 @@ CompileOptions prepareOptions() compileOptions.appendGCtorsToMethod = appendGCtorsToMethod.getValue(); compileOptions.strictNullChecks = strictNullChecks.getValue(); compileOptions.enableFastMath = enableFastMath.getValue(); + // Not for a program read from stdin, which has no folder to keep it in, nor for + // --dump-object-file, which writes the object the JIT makes of the one module. + compileOptions.jitCache = compileOptions.isJit && jitCache.getValue() && inputFilename != "-" && !dumpObjectFile; if (!outputFilename.empty()) { diff --git a/tslang/tslang/tslang.cpp b/tslang/tslang/tslang.cpp index b89797c81..66806a90e 100644 --- a/tslang/tslang/tslang.cpp +++ b/tslang/tslang/tslang.cpp @@ -67,6 +67,7 @@ int dumpObjOrAssembly(int, char **, mlir::ModuleOp, CompileOptions&); int declarationInline(int, char **, mlir::MLIRContext &, llvm::SourceMgr &, llvm::StringRef, CompileOptions&, std::string&); int buildExe(int, char **, std::string, std::string, CompileOptions&); int runJit(int, char **, mlir::ModuleOp, CompileOptions&); +int runJitCached(int, char **, CompileOptions&); extern cl::OptionCategory ObjOrAssemblyCategory; cl::OptionCategory TypeScriptCompilerCategory("Compiler Options"); @@ -117,6 +118,9 @@ cl::opt dumpObjectFile{"dump-object-file", cl::Hidden, cl::desc("Dump JITt cl::opt objectFilename{"object-filename", cl::Hidden, cl::desc("Dump JITted-compiled object to file .o"), cl::cat(TypeScriptCompilerDebugCategory)}; +cl::opt jitCache{"jit-cache", cl::desc("--emit=jit: compile the program and each .ts module it imports into an object of its own, kept in a cache folder and loaded from there while its sources do not change (on by default; --jit-cache=false compiles everything into one module every run)"), cl::init(true), cl::cat(TypeScriptCompilerCategory)}; +cl::opt jitCacheDir{"jit-cache-dir", cl::desc("--emit=jit: the folder the JIT cache keeps its objects in (default: '__jit' beside each source file)"), cl::value_desc("folder"), cl::cat(TypeScriptCompilerCategory)}; + cl::opt verbose{"verbose", cl::Hidden, cl::desc("Verbose output"), cl::init(false), cl::cat(TypeScriptCompilerDebugCategory)}; cl::opt printOp{"print-op", cl::Hidden, cl::desc("Print Op on Diagnostic"), cl::init(false), cl::cat(TypeScriptCompilerDebugCategory)}; cl::opt printStackTrace{"print-stack-trace", cl::Hidden, cl::desc("Print stack trace on Diagnostic"), cl::init(false), cl::cat(TypeScriptCompilerDebugCategory)}; @@ -295,6 +299,31 @@ static int ReportWarningWithoutBreak(int reportType, char *message, int *returnV } #endif +// A context with every dialect the compiler generates and lowers through loaded. +std::unique_ptr createMLIRContext() +{ + mlir::DialectRegistry registry; + registerAllExtensions(registry); + + auto mlirContext = std::make_unique(registry); + // Load our Dialect in this MLIR Context. + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); + mlirContext->getOrLoadDialect(); +#ifdef ENABLE_ASYNC + mlirContext->getOrLoadDialect(); +#endif + + mlirContext->printOpOnDiagnostic(printOp.getValue()); + mlirContext->printStackTraceOnDiagnostic(printStackTrace.getValue()); + return mlirContext; +} + int main(int argc, char **argv) { #if _MSC_VER && _DEBUG @@ -412,27 +441,6 @@ int main(int argc, char **argv) return dumpAST(); } - // If we aren't dumping the AST, then we are compiling with/to MLIR. - mlir::DialectRegistry registry; - registerAllExtensions(registry); - - mlir::MLIRContext mlirContext(registry); - // Load our Dialect in this MLIR Context. - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); - mlirContext.getOrLoadDialect(); -#ifdef ENABLE_ASYNC - mlirContext.getOrLoadDialect(); -#endif - - mlirContext.printOpOnDiagnostic(printOp.getValue()); - mlirContext.printStackTraceOnDiagnostic(printStackTrace.getValue()); - auto compileOptions = prepareOptions(); if (!prepareDefaultLib(compileOptions)) @@ -440,6 +448,16 @@ int main(int argc, char **argv) return 0; } + // The JIT with its cache compiles each file it needs by itself, and only if it has to. + if (emitAction == Action::RunJIT && compileOptions.jitCache) + { + return runJitCached(argc, argv, compileOptions); + } + + // If we aren't dumping the AST, then we are compiling with/to MLIR. + auto mlirContextPtr = createMLIRContext(); + auto &mlirContext = *mlirContextPtr; + llvm::SourceMgr sourceMgr; mlir::OwningOpRef module; if (auto error = compileTypeScriptFileIntoMLIR(mlirContext, sourceMgr, module, compileOptions)) From cb34bc73e91de6e21cd16f3964991d0ac22ce7e0 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 29 Sep 2026 16:59:53 +0000 Subject: [PATCH 2/2] jit: (std::max) so windows.h's max macro does not break the Win64 build Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VPAHhg8c9NEdoYNw52AtTg --- tslang/tslang/jit.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tslang/tslang/jit.cpp b/tslang/tslang/jit.cpp index 4e1b5c0cf..b96263707 100644 --- a/tslang/tslang/jit.cpp +++ b/tslang/tslang/jit.cpp @@ -276,7 +276,7 @@ class JitSectionMemoryManager : public llvm::SectionMemoryManager } image->base = imageBase; - image->end = std::max(image->end, reinterpret_cast(addr) + size); + image->end = (std::max)(image->end, reinterpret_cast(addr) + size); #endif }