System image compression with zstd (#59227)

Revived version of #48244, with a slightly different approach. This
version looks for a function pointer called `jl_image_unpack` inside
compiled system images and invokes it to get the `jl_image_buf_t`
struct. Two implementations, `jl_image_unpack_zstd` and
`jl_image_unpack_uncomp` are provided (for comparison). The zstd
compression is applied only to the heap image, and not the compiled
code, since that can be shared across Julia processes.

TODO: test a few different compression settings and enable by default.

Example data from un-trimmed juliac "hello world":
```
156M  hello-uncomp
 43M  hello-zstd
 48M  hello-zstd-1
 45M  hello-zstd-5
 43M  hello-zstd-15
 39M  hello-zstd-22

$ hyperfine -w3 ./hello-uncomp 
Benchmark 1: ./hello-uncomp
  Time (mean ± σ):      74.4 ms ±   0.8 ms    [User: 51.9 ms, System: 19.0 ms]
  Range (min … max):    73.0 ms …  76.6 ms    39 runs

$ hyperfine -w3 ./hello-zstd-1
Benchmark 1: ./hello-zstd-1
  Time (mean ± σ):     152.4 ms ±   0.5 ms    [User: 138.2 ms, System: 12.0 ms]
  Range (min … max):   151.4 ms … 153.2 ms    19 runs
 
$ hyperfine -w3 ./hello-zstd-5 
Benchmark 1: ./hello-zstd-5
  Time (mean ± σ):     154.3 ms ±   0.5 ms    [User: 139.6 ms, System: 12.4 ms]
  Range (min … max):   153.5 ms … 155.2 ms    19 runs

$ hyperfine -w3 ./hello-zstd-15
Benchmark 1: ./hello-zstd-15
  Time (mean ± σ):     135.9 ms ±   0.5 ms    [User: 121.6 ms, System: 12.0 ms]
  Range (min … max):   135.1 ms … 136.5 ms    21 runs
 
$ hyperfine -w3 ./hello-zstd-22
Benchmark 1: ./hello-zstd-22
  Time (mean ± σ):     149.0 ms ±   0.6 ms    [User: 134.7 ms, System: 12.1 ms]
  Range (min … max):   147.7 ms … 150.4 ms    19 runs
```

---------

Co-authored-by: Gabriel Baraldi <[email protected]>

Author: xal-0

Make.inc

@@ -1591,7 +1591,7 @@ JLDFLAGS += -Wl,--stack,8388608
 ifeq ($(ARCH),i686)
 JLDFLAGS += -Wl,--large-address-aware
 endif
-JCPPFLAGS += -D_WIN32_WINNT=0x0502
+JCPPFLAGS += -D_WIN32_WINNT=_WIN32_WINNT_WIN8
 UNTRUSTED_SYSTEM_LIBM := 1
 # Use hard links for files on windows, rather than soft links
 #   https://stackoverflow.com/questions/3648819/how-to-make-a-symbolic-link-with-cygwin-in-windows-7

base/options.jl

@@ -67,6 +67,7 @@ struct JLOptions
     task_metrics::Int8
     timeout_for_safepoint_straggler_s::Int16
     gc_sweep_always_full::Int8
+    compress_sysimage::Int8
 end
 
 # This runs early in the sysimage != is not defined yet

base/util.jl

@@ -245,6 +245,9 @@ function julia_cmd(julia=joinpath(Sys.BINDIR, julia_exename()); cpu_target::Unio
     if opts.use_sysimage_native_code == 0
         push!(addflags, "--sysimage-native-code=no")
     end
+    if opts.compress_sysimage == 1
+        push!(addflags, "--compress-sysimage=yes")
+    end
     return `$julia -C $cpu_target -J$image_file $addflags`
 end
 

src/Makefile

@@ -524,7 +524,7 @@ $(build_shlibdir)/lib%Plugin.$(SHLIB_EXT): $(SRCDIR)/clangsa/%.cpp $(LLVM_CONFIG
 # before attempting this static analysis, so that all necessary headers
 # and dependencies are properly installed:
 #   make -C src install-analysis-deps
-ANALYSIS_DEPS := llvm clang llvm-tools libuv utf8proc
+ANALYSIS_DEPS := llvm clang llvm-tools libuv utf8proc zstd
 ifeq ($(OS),Darwin)
 ANALYSIS_DEPS += llvmunwind
 else ifeq ($(OS),OpenBSD)

src/aotcompile.cpp

@@ -42,9 +42,10 @@
 #include <llvm/Support/FormatAdapters.h>
 #include <llvm/Linker/Linker.h>
 
-
 using namespace llvm;
 
+#include <zstd.h>
+
 #include "jitlayers.h"
 #include "serialize.h"
 #include "julia_assert.h"
@@ -2148,27 +2149,53 @@ void jl_dump_native_impl(void *native_code,
         sysimgM.setDataLayout(DL);
         sysimgM.setStackProtectorGuard(StackProtectorGuard);
         sysimgM.setOverrideStackAlignment(OverrideStackAlignment);
-        Constant *data = ConstantDataArray::get(Context,
-            ArrayRef<uint8_t>((const unsigned char*)z->buf, z->size));
+
+        int compression = jl_options.compress_sysimage ? 15 : 0;
+        ArrayRef<char> sysimg_data{z->buf, (size_t)z->size};
+        SmallVector<char, 0> compressed_data;
+        if (compression) {
+            compressed_data.resize(ZSTD_compressBound(z->size));
+            size_t comp_size = ZSTD_compress(compressed_data.data(), compressed_data.size(),
+                                             z->buf, z->size, compression);
+            compressed_data.resize(comp_size);
+            sysimg_data = compressed_data;
+            ios_close(z);
+            free(z);
+        }
+
+        Constant *data = ConstantDataArray::get(Context, sysimg_data);
         auto sysdata = new GlobalVariable(sysimgM, data->getType(), false,
                                      GlobalVariable::ExternalLinkage,
                                      data, "jl_system_image_data");
-        sysdata->setAlignment(Align(64));
+        sysdata->setAlignment(Align(jl_page_size));
 #if JL_LLVM_VERSION >= 180000
         sysdata->setCodeModel(CodeModel::Large);
 #else
         if (TheTriple.isX86() && TheTriple.isArch64Bit() && TheTriple.isOSLinux())
             sysdata->setSection(".ldata");
 #endif
         addComdat(sysdata, TheTriple);
-        Constant *len = ConstantInt::get(sysimgM.getDataLayout().getIntPtrType(Context), z->size);
+        Constant *len = ConstantInt::get(sysimgM.getDataLayout().getIntPtrType(Context), sysimg_data.size());
         addComdat(new GlobalVariable(sysimgM, len->getType(), true,
                                      GlobalVariable::ExternalLinkage,
                                      len, "jl_system_image_size"), TheTriple);
-        // Free z here, since we've copied out everything into data
-        // Results in serious memory savings
-        ios_close(z);
-        free(z);
+
+        const char *unpack_func = compression ? "jl_image_unpack_zstd" : "jl_image_unpack_uncomp";
+        auto unpack = new GlobalVariable(sysimgM, DL.getIntPtrType(Context), true,
+                                         GlobalVariable::ExternalLinkage, nullptr,
+                                         unpack_func);
+        addComdat(new GlobalVariable(sysimgM, PointerType::getUnqual(Context), true,
+                                     GlobalVariable::ExternalLinkage, unpack,
+                                     "jl_image_unpack"),
+                  TheTriple);
+
+        if (!compression) {
+            // Free z here, since we've copied out everything into data
+            // Results in serious memory savings
+            ios_close(z);
+            free(z);
+        }
+        compressed_data.clear();
         // Note that we don't set z to null, this allows the check in WRITE_ARCHIVE
         // to function as expected
         // no need to free the module/context, destructor handles that

src/jloptions.c

@@ -160,6 +160,7 @@ JL_DLLEXPORT void jl_init_options(void)
                         0, // task_metrics
                         -1, // timeout_for_safepoint_straggler_s
                         0, // gc_sweep_always_full
+                        0, // compress_sysimage
     };
     jl_options_initialized = 1;
 }
@@ -311,7 +312,10 @@ static const char opts_hidden[]  =
     " --strip-metadata                              Remove docstrings and source location info from\n"
     "                                               system image\n"
     " --strip-ir                                    Remove IR (intermediate representation) of compiled\n"
-    "                                               functions\n\n"
+    "                                               functions\n"
+    " --compress-sysimage={yes|no*}                 Compress the sys/pkgimage heap at the expense of\n"
+    "                                               slightly increased load time.\n"
+    "\n"
 
     // compiler debugging and experimental (see the devdocs for tips on using these options)
     " --experimental                                Enable the use of experimental (alpha) features\n"
@@ -407,6 +411,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
            opt_permalloc_pkgimg,
            opt_trim,
            opt_experimental_features,
+           opt_compress_sysimage,
     };
     static const char* const shortopts = "+vhqH:e:E:L:J:C:it:p:O:g:m:";
     static const struct option longopts[] = {
@@ -478,6 +483,7 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
         { "heap-target-increment", required_argument, 0, opt_heap_target_increment },
         { "gc-sweep-always-full", no_argument, 0, opt_gc_sweep_always_full },
         { "trim",  optional_argument, 0, opt_trim },
+        { "compress-sysimage", required_argument, 0, opt_compress_sysimage },
         { 0, 0, 0, 0 }
     };
 
@@ -1060,6 +1066,12 @@ JL_DLLEXPORT void jl_parse_opts(int *argcp, char ***argvp)
             else
                 jl_errorf("julia: invalid argument to --task-metrics={yes|no} (%s)", optarg);
             break;
+        case opt_compress_sysimage:
+            if (!strcmp(optarg,"yes"))
+                jl_options.compress_sysimage = 1;
+            else if (!strcmp(optarg,"no"))
+                jl_options.compress_sysimage = 0;
+            break;
         default:
             jl_errorf("julia: unhandled option -- %c\n"
                       "This is a bug, please report it.", c);

src/jloptions.h

@@ -71,6 +71,7 @@ typedef struct {
     int8_t task_metrics;
     int16_t timeout_for_safepoint_straggler_s;
     int8_t gc_sweep_always_full;
+    int8_t compress_sysimage;
 } jl_options_t;
 
 #endif

src/staticdata.c

@@ -72,15 +72,20 @@ External links:
 #include <stdio.h> // printf
 #include <inttypes.h> // PRIxPTR
 
+#include <zstd.h>
+
 #include "julia.h"
 #include "julia_internal.h"
 #include "julia_gcext.h"
 #include "builtin_proto.h"
 #include "processor.h"
 #include "serialize.h"
 
-#ifndef _OS_WINDOWS_
+#ifdef _OS_WINDOWS_
+#include <memoryapi.h>
+#else
 #include <dlfcn.h>
+#include <sys/mman.h>
 #endif
 
 #include "valgrind.h"
@@ -3630,14 +3635,75 @@ JL_DLLEXPORT jl_image_buf_t jl_preload_sysimg(const char *fname)
     }
 }
 
-// From a shared library handle, verify consistency and return a jl_image_buf_t
-static jl_image_buf_t get_image_buf(void *handle, int is_pkgimage)
+typedef void jl_image_unpack_func_t(void *handle, jl_image_buf_t *image);
+
+static void jl_prefetch_system_image(const char *data, size_t size)
+{
+    size_t page_size = jl_getpagesize(); /* jl_page_size is not set yet when loading sysimg */
+    void *start = (void *)((uintptr_t)data & ~(page_size - 1));
+    size_t size_aligned = LLT_ALIGN(size, page_size);
+#ifdef _OS_WINDOWS_
+    WIN32_MEMORY_RANGE_ENTRY entry = {start, size_aligned};
+    PrefetchVirtualMemory(GetCurrentProcess(), 1, &entry, 0);
+#else
+    madvise(start, size_aligned, MADV_WILLNEED);
+#endif
+}
+
+JL_DLLEXPORT void jl_image_unpack_uncomp(void *handle, jl_image_buf_t *image)
+{
+    size_t *plen;
+    jl_dlsym(handle, "jl_system_image_size", (void **)&plen, 1);
+    jl_dlsym(handle, "jl_system_image_data", (void **)&image->data, 1);
+    jl_dlsym(handle, "jl_image_pointers", (void**)&image->pointers, 1);
+    image->size = *plen;
+    jl_prefetch_system_image(image->data, image->size);
+}
+
+JL_DLLEXPORT void jl_image_unpack_zstd(void *handle, jl_image_buf_t *image)
 {
     size_t *plen;
     const char *data;
-    const void *pointers;
-    uint64_t base;
+    jl_dlsym(handle, "jl_system_image_size", (void **)&plen, 1);
+    jl_dlsym(handle, "jl_system_image_data", (void **)&data, 1);
+    jl_dlsym(handle, "jl_image_pointers", (void **)&image->pointers, 1);
+    jl_prefetch_system_image(data, *plen);
+    image->size = ZSTD_getFrameContentSize(data, *plen);
+    size_t page_size = jl_getpagesize(); /* jl_page_size is not set yet when loading sysimg */
+    size_t aligned_size = LLT_ALIGN(image->size, page_size);
+#if defined(_OS_WINDOWS_)
+    size_t large_page_size = GetLargePageMinimum();
+    if (image->size > 4 * large_page_size) {
+        size_t aligned_size = LLT_ALIGN(image->size, large_page_size);
+        image->data = (char *)VirtualAlloc(
+            NULL, aligned_size, MEM_COMMIT | MEM_RESERVE | MEM_LARGE_PAGES, PAGE_READWRITE);
+    }
+    else {
+        image->data = (char *)VirtualAlloc(NULL, aligned_size, MEM_COMMIT | MEM_RESERVE,
+                                           PAGE_READWRITE);
+    }
+#else
+    image->data =
+        (char *)mmap(NULL, aligned_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+#endif
+    if (!image->data || image->data == (void *)-1) {
+        jl_printf(JL_STDERR, "ERROR: failed to allocate space for system image\n");
+        jl_exit(1);
+    }
+
+    ZSTD_decompress((void *)image->data, image->size, data, *plen);
+    size_t len = (*plen) & ~(page_size - 1);
+#ifdef _OS_WINDOWS_
+    if (len)
+        VirtualFree((void *)data, len, MEM_RELEASE);
+#else
+    munmap((void *)data, len);
+#endif
+}
 
+// From a shared library handle, verify consistency and return a jl_image_buf_t
+static jl_image_buf_t get_image_buf(void *handle, int is_pkgimage)
+{
     // verify that the linker resolved the symbols in this image against ourselves (libjulia-internal)
     void** (*get_jl_RTLD_DEFAULT_handle_addr)(void) = NULL;
     if (handle != jl_RTLD_DEFAULT_handle) {
@@ -3646,38 +3712,41 @@ static jl_image_buf_t get_image_buf(void *handle, int is_pkgimage)
             jl_error("Image file failed consistency check: maybe opened the wrong version?");
     }
 
+    jl_image_unpack_func_t **unpack;
+    jl_image_buf_t image = {
+        .kind = JL_IMAGE_KIND_SO,
+        .pointers = NULL,
+        .data = NULL,
+        .size = 0,
+        .base = 0,
+    };
+
     // verification passed, lookup the buffer pointers
     if (jl_system_image_size == 0 || is_pkgimage) {
         // in the usual case, the sysimage was not statically linked to libjulia-internal
         // look up the external sysimage symbols via the dynamic linker
-        jl_dlsym(handle, "jl_system_image_size", (void **)&plen, 1);
-        jl_dlsym(handle, "jl_system_image_data", (void **)&data, 1);
-        jl_dlsym(handle, "jl_image_pointers", (void**)&pointers, 1);
-    } else {
+        jl_dlsym(handle, "jl_image_unpack", (void **)&unpack, 1);
+        (*unpack)(handle, &image);
+    }
+    else {
         // the sysimage was statically linked directly against libjulia-internal
         // use the internal symbols
-        plen = &jl_system_image_size;
-        pointers = &jl_image_pointers;
-        data = &jl_system_image_data;
+        image.size = jl_system_image_size;
+        image.pointers = &jl_image_pointers;
+        image.data = &jl_system_image_data;
     }
 
 #ifdef _OS_WINDOWS_
-    base = (intptr_t)handle;
+    image.base = (intptr_t)handle;
 #else
     Dl_info dlinfo;
-    if (dladdr((void*)pointers, &dlinfo) != 0)
-        base = (intptr_t)dlinfo.dli_fbase;
+    if (dladdr((void*)image.pointers, &dlinfo) != 0)
+        image.base = (intptr_t)dlinfo.dli_fbase;
     else
-        base = 0;
+        image.base = 0;
 #endif
 
-    return (jl_image_buf_t) {
-        .kind = JL_IMAGE_KIND_SO,
-        .pointers = pointers,
-        .data = data,
-        .size = *plen,
-        .base = base,
-    };
+    return image;
 }
 
 // Allow passing in a module handle directly, rather than a path