JIT cache: each .ts file is compiled into an object kept in __jit

.gitignore

@@ -12,3 +12,6 @@ package-lock.json
 node_modules
 gc-*
 libatomic_ops-*
+
+# the JIT cache (tslang --emit=jit)
+__jit/

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.<hash>.o        the object
+__jit/hello.ts.<hash>.o.deps   what it was compiled from
+```
+
+- `--jit-cache-dir=<folder>` 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,

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<std::string> sourceImports;
+    std::vector<std::string> 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.

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<bool> 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);

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=$<TARGET_FILE:tslang>"
+                 "-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)

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(<expected output>) - 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")

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";
+    }
+}

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.");

tslang/test/tester/jit-cache/cycle_module.ts

@@ -0,0 +1,5 @@
+import './cycle_main'
+
+export function fromModule() {
+    return "module+" + fromMain();
+}

tslang/test/tester/jit-cache/greeting.ts

@@ -0,0 +1,3 @@
+import './greeting_module'
+
+print(greeting());

tslang/test/tester/jit-cache/greeting_module.ts

@@ -0,0 +1,3 @@
+export function greeting() {
+    return "one";
+}

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 = <string>e;
+}
+
+assert(caught == "boom");
+
+print("done.");

tslang/test/tester/jit-cache/other_module.ts

@@ -0,0 +1,5 @@
+import './base_module'
+
+export function bumpFromOther() {
+    return bump();
+}

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;

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