Write JIT pack file using same method as native code gen

Fixes #1159

Author: nickg

NEWS.md

@@ -32,7 +32,7 @@
   of a conversion function (#1171).
 - Type marks are now allowed in the prefix of selected names in relaxed
   mode (#1173).
-- Several other minor bugs were resolved (#1174).
+- Several other minor bugs were resolved (#1174, #1159).
 
 ## Version 1.15.2 - 2025-03-01
 - Fixed invalid LLVM IR generation which could cause a crash with LLVM

src/Makemodule.am

@@ -54,16 +54,13 @@ lib_libnvc_a_SOURCES = \
 	src/driver.h \
 	src/driver.c \
 	src/inst.h \
-	src/inst.c
+	src/inst.c \
+	src/cgen.c
 
 if ENABLE_SERVER
 lib_libnvc_a_SOURCES += src/server.c src/server.h
 endif
 
-if ENABLE_LLVM
-lib_libnvc_a_SOURCES += src/cgen.c
-endif
-
 bin_nvc_SOURCES = src/nvc.c
 
 EXTRA_bin_nvc_DEPENDENCIES = src/symbols.txt

src/cgen.c

@@ -56,8 +56,10 @@ typedef struct {
    unit_registry_t *registry;
 } discover_args_t;
 
+#ifdef ENABLE_LLVM
 static A(char *) link_args;
 static A(char *) cleanup_files = AINIT;
+#endif
 
 #define UNITS_PER_JOB 25
 
@@ -163,6 +165,7 @@ static void cgen_walk_hier(unit_list_t *units, hset_t *seen, tree_t block,
    }
 }
 
+#ifdef ENABLE_LLVM
 static void cleanup_temp_dll(void)
 {
    for (int i = 0; i < cleanup_files.count; i++) {
@@ -305,7 +308,7 @@ static void cgen_async_work(void *context, void *arg)
 
 static void cgen_partition_jobs(unit_list_t *units, workq_t *wq,
                                 const char *base_name, int units_per_job,
-                                tree_t top, obj_list_t *objs)
+                                obj_list_t *objs)
 {
    int counter = 0;
 
@@ -337,7 +340,44 @@ static void cgen_partition_jobs(unit_list_t *units, workq_t *wq,
    }
 }
 
-void cgen(tree_t top, unit_registry_t *ur, mir_context_t *mc, jit_t *jit)
+static void cgen_native(ident_t name, jit_t *jit, unit_list_t *units)
+{
+   workq_t *wq = workq_new(jit);
+
+   obj_list_t objs = AINIT;
+   cgen_partition_jobs(units, wq, istr(name), UNITS_PER_JOB, &objs);
+
+   workq_start(wq);
+   workq_drain(wq);
+
+   progress("code generation for %d units", units->count);
+
+   cgen_link(istr(name), objs.items, objs.count);
+
+   for (unsigned i = 0; i < objs.count; i++)
+      free(objs.items[i]);
+   ACLEAR(objs);
+
+   workq_free(wq);
+}
+#endif   // ENABLE_LLVM
+
+static void cgen_jit_pack(ident_t name, jit_t *jit, unit_list_t *units)
+{
+   const char *fname LOCAL = xasprintf("_%s.pack", istr(name));
+
+   FILE *f = lib_fopen(lib_work(), fname, "wb");
+   if (f == NULL)
+      fatal_errno("fopen: %s", fname);
+
+   jit_write_pack(jit, units->items, units->count, f);
+   fclose(f);
+
+   progress("writing JIT pack");
+}
+
+void cgen(tree_t top, unit_registry_t *ur, mir_context_t *mc, jit_t *jit,
+          cgen_mode_t mode)
 {
    assert(tree_kind(top) == T_ELAB);
 
@@ -359,28 +399,23 @@ void cgen(tree_t top, unit_registry_t *ur, mir_context_t *mc, jit_t *jit)
    hset_free(seen);
    seen = NULL;
 
-   workq_t *wq = workq_new(jit);
-
-   ident_t name = tree_ident(top);
-
-   obj_list_t objs = AINIT;
-   cgen_partition_jobs(&units, wq, istr(name), UNITS_PER_JOB, top, &objs);
-
-   workq_start(wq);
-   workq_drain(wq);
-
-   progress("code generation for %d units", units.count);
-
-   cgen_link(istr(name), objs.items, objs.count);
-
-   for (unsigned i = 0; i < objs.count; i++)
-      free(objs.items[i]);
-   ACLEAR(objs);
+   switch (mode) {
+   case CGEN_NATIVE:
+#ifdef ENABLE_LLVM
+      cgen_native(tree_ident(top), jit, &units);
+#else
+      fatal("native code generation not enabled in this build");
+#endif
+      break;
+   case CGEN_JIT_PACK:
+      cgen_jit_pack(tree_ident(top), jit, &units);
+      break;
+   }
 
    ACLEAR(units);
-   workq_free(wq);
 }
 
+#ifdef ENABLE_LLVM
 static void preload_walk_index(lib_t lib, ident_t ident, int kind, void *ctx)
 {
    discover_args_t *args = ctx;
@@ -578,3 +613,4 @@ void aotgen(const char *outfile, char **argv, int argc)
    unit_registry_free(ur);
    mir_context_free(mc);
 }
+#endif   // ENABLE_LLVM

src/jit/jit-pack.c

@@ -23,7 +23,6 @@
 #include "jit/jit-priv.h"
 #include "lib.h"
 #include "object.h"
-#include "vcode.h"
 
 #include <assert.h>
 #include <stdlib.h>
@@ -618,47 +617,39 @@ static void write_fully(const void *buf, size_t size, FILE *f)
       fatal_errno("fwrite");
 }
 
-static void write_children(jit_t *j, vcode_unit_t vu, pack_writer_t *pw,
-                           FILE *file)
+void jit_write_pack(jit_t *j, const ident_t *units, size_t count, FILE *f)
 {
-   ident_t ident = vcode_unit_name(vu);
-   jit_handle_t handle = jit_compile(j, ident);
+   pack_writer_t *pw = pack_writer_new();
 
-   uint8_t bytes[10];
+   pack_header_t header = {};
+   write_fully(&header, sizeof(header), f);
 
-   const uint32_t name = pack_writer_get_string(pw, istr(ident));
-   const int name_nbytes = encode_number(name, bytes);
-   write_fully(bytes, name_nbytes, file);
+   for (size_t i = 0; i < count; i++) {
+      jit_handle_t handle = jit_compile(j, units[i]);
 
-   uint8_t *buf;
-   size_t size;
-   pack_writer_emit(pw, j, handle, &buf, &size);
+      uint8_t bytes[10];
 
-   const int size_nbytes = encode_number(size, bytes);
-   write_fully(bytes, size_nbytes, file);
+      const uint32_t name = pack_writer_get_string(pw, istr(units[i]));
+      const int name_nbytes = encode_number(name, bytes);
+      write_fully(bytes, name_nbytes, f);
 
-   jit_func_t *f = jit_get_func(j, handle);
+      uint8_t *buf;
+      size_t size;
+      pack_writer_emit(pw, j, handle, &buf, &size);
 
-   const int cpool_nbytes = encode_number(f->cpoolsz, bytes);
-   write_fully(bytes, cpool_nbytes, file);
+      const int size_nbytes = encode_number(size, bytes);
+      write_fully(bytes, size_nbytes, f);
 
-   write_fully(buf, size, file);
+      jit_func_t *func = jit_get_func(j, handle);
 
-   if (f->cpoolsz > 0)
-      write_fully(f->cpool, f->cpoolsz, file);
+      const int cpool_nbytes = encode_number(func->cpoolsz, bytes);
+      write_fully(bytes, cpool_nbytes, f);
 
-   for (vcode_unit_t it = vcode_unit_child(vu); it; it = vcode_unit_next(it))
-      write_children(j, it, pw, file);
-}
+      write_fully(buf, size, f);
 
-void jit_write_pack(jit_t *j, vcode_unit_t root, FILE *f)
-{
-   pack_writer_t *pw = pack_writer_new();
-
-   pack_header_t header = {};
-   write_fully(&header, sizeof(header), f);
-
-   write_children(j, root, pw, f);
+      if (func->cpoolsz > 0)
+         write_fully(func->cpool, func->cpoolsz, f);
+   }
 
    memcpy(header.magic, PACK_MAGIC, sizeof(header.magic));
    header.strtab = ftell(f);

src/jit/jit.h

@@ -96,7 +96,7 @@ tlab_t jit_null_tlab(jit_t *j);
 jit_pack_t *jit_pack_new(void);
 void jit_pack_free(jit_pack_t *jp);
 
-void jit_write_pack(jit_t *j, vcode_unit_t root, FILE *f);
+void jit_write_pack(jit_t *j, const ident_t *units, size_t count, FILE *f);
 jit_pack_t *jit_read_pack(FILE *f);
 
 __attribute__((format(printf, 3, 4)))

src/nvc.c

@@ -393,6 +393,16 @@ static int parse_optimise_level(const char *str)
    return level;
 }
 
+static void load_jit_pack(jit_t *jit, tree_t top)
+{
+   char *name LOCAL = xasprintf("_%s.pack", istr(tree_ident(top)));
+   FILE *f = lib_fopen(lib_work(), name, "rb");
+   if (f != NULL) {
+      jit_load_pack(jit, f);
+      fclose(f);
+   }
+}
+
 static int elaborate(int argc, char **argv, cmd_state_t *state)
 {
    static struct option long_options[] = {
@@ -556,8 +566,9 @@ static int elaborate(int argc, char **argv, cmd_state_t *state)
    if (error_count() > 0)
       return EXIT_FAILURE;
 
-   char *pack_name LOCAL = xasprintf("_%s.pack", istr(state->top_level));
-   char *dll_name LOCAL = xasprintf("_%s." DLL_EXT, istr(tree_ident(top)));
+   const char *elab_name = istr(tree_ident(top));
+   char *pack_name LOCAL = xasprintf("_%s.pack", elab_name);
+   char *dll_name LOCAL = xasprintf("_%s." DLL_EXT, elab_name);
 
    // Delete any existing generated code to avoid accidentally loading
    // the wrong version later
@@ -569,25 +580,10 @@ static int elaborate(int argc, char **argv, cmd_state_t *state)
       progress("saving library");
    }
 
-   if (use_jit && !no_save) {
-      FILE *f = lib_fopen(state->work, pack_name, "wb");
-      if (f == NULL)
-         fatal_errno("fopen: %s", pack_name);
-
-      ident_t b0 = tree_ident(tree_stmt(top, 0));
-      ident_t root = ident_prefix(lib_name(state->work), b0, '.');
-
-      vcode_unit_t vu = unit_registry_get(state->registry, root);
-      assert(vu != NULL);
-
-      jit_write_pack(state->jit, vu, f);
-      fclose(f);
-
-      progress("writing JIT pack");
-   }
-
    if (!use_jit)
-      LLVM_ONLY(cgen(top, state->registry, state->mir, state->jit));
+      cgen(top, state->registry, state->mir, state->jit, CGEN_NATIVE);
+   else if (!no_save)
+      cgen(top, state->registry, state->mir, state->jit, CGEN_JIT_PACK);
 
    if (!use_jit || cover != NULL) {
       // Must discard current JIT state to load AOT library later
@@ -949,12 +945,7 @@ static int run_cmd(int argc, char **argv, cmd_state_t *state)
    jit_load_dll(state->jit, tree_ident(top));
 #endif
 
-   char *name LOCAL = xasprintf("_%s.pack", istr(state->top_level));
-   FILE *f = lib_fopen(state->work, name, "rb");
-   if (f != NULL) {
-      jit_load_pack(state->jit, f);
-      fclose(f);
-   }
+   load_jit_pack(state->jit, top);
 
    if (state->cover == NULL)
       state->cover = load_coverage(meta, state->jit);
@@ -1483,12 +1474,7 @@ static int do_cmd(int argc, char **argv, cmd_state_t *state)
       jit_load_dll(state->jit, tree_ident(top));
 #endif
 
-      char *name LOCAL = xasprintf("_%s.pack", istr(state->top_level));
-      FILE *f = lib_fopen(state->work, name, "rb");
-      if (f != NULL) {
-         jit_load_pack(state->jit, f);
-         fclose(f);
-      }
+      load_jit_pack(state->jit, top);
 
       shell_reset(sh, top);
    }
@@ -1566,12 +1552,7 @@ static int interact_cmd(int argc, char **argv, cmd_state_t *state)
       jit_load_dll(state->jit, tree_ident(top));
 #endif
 
-      char *name LOCAL = xasprintf("_%s.pack", istr(state->top_level));
-      FILE *f = lib_fopen(state->work, name, "rb");
-      if (f != NULL) {
-         jit_load_pack(state->jit, f);
-         fclose(f);
-      }
+      load_jit_pack(state->jit, top);
 
       shell_reset(sh, top);
    }

src/phase.h

@@ -44,8 +44,14 @@ tree_t elab(object_t *top, jit_t *jit, unit_registry_t *ur, mir_context_t *mc,
 // Set the value of a top-level generic
 void elab_set_generic(const char *name, const char *value);
 
+typedef enum {
+   CGEN_NATIVE,
+   CGEN_JIT_PACK,
+} cgen_mode_t;
+
 // Generate LLVM bitcode for a design unit
-void cgen(tree_t top, unit_registry_t *ur, mir_context_t *mc, jit_t *jit);
+void cgen(tree_t top, unit_registry_t *ur, mir_context_t *mc, jit_t *jit,
+          cgen_mode_t mode);
 
 // Generate ahead-of-time preload library
 void aotgen(const char *outfile, char **argv, int argc);