JIT: Add flushing of stream and backend

Signed-off-by: Paul Guyot <pguyot@kallisys.net>

Author: pguyot

libs/estdlib/src/code_server.erl

@@ -164,8 +164,8 @@ load(Module) ->
                             code_server:literal_resolver(Module, Index)
                         end,
                         TypeResolver = fun(Index) -> code_server:type_resolver(Module, Index) end,
-                        Stream0 = jit:stream(jit_mmap_size(byte_size(Code))),
-                        {BackendModule, BackendState0} = jit:backend(Stream0),
+                        {StreamModule, Stream0} = jit:stream(jit_mmap_size(byte_size(Code))),
+                        {BackendModule, BackendState0} = jit:backend(StreamModule, Stream0),
                         {LabelsCount, BackendState1} = jit:compile(
                             Code,
                             AtomResolver,
@@ -174,12 +174,12 @@ load(Module) ->
                             BackendModule,
                             BackendState0
                         ),
-                        BackendState2 = BackendModule:flush(BackendState1),
-                        Stream1 = BackendModule:stream(BackendState2),
-                        code_server:set_native_code(Module, LabelsCount, Stream1),
+                        Stream1 = BackendModule:stream(BackendState1),
+                        Stream2 = StreamModule:flush(Stream1),
+                        code_server:set_native_code(Module, LabelsCount, Stream2),
                         End = erlang:system_time(millisecond),
                         io:format("~B ms (bytecode: ~B bytes, native code: ~B bytes)\n", [
-                            End - Start, byte_size(Code), BackendModule:offset(BackendState2)
+                            End - Start, byte_size(Code), BackendModule:offset(BackendState1)
                         ])
                     catch
                         T:V:S ->

libs/jit/src/jit.erl

@@ -22,7 +22,7 @@
 
 -export([
     stream/1,
-    backend/1,
+    backend/2,
     beam_chunk_header/3,
     compile/6
 ]).
@@ -148,7 +148,8 @@ compile(
     },
     {State1, MSt2} = first_pass(Opcodes, MMod, MSt1, State0),
     MSt3 = second_pass(MMod, MSt2, State1),
-    {LabelsCount, MSt3};
+    MSt4 = MMod:flush(MSt3),
+    {LabelsCount, MSt4};
 compile(
     <<16:32, 0:32, OpcodeMax:32, _LabelsCount:32, _FunctionsCount:32, _Opcodes/binary>>,
     _AtomResolver,
@@ -3792,7 +3793,7 @@ variant() ->
 
 %% @doc Instantiate backend for this platform
 %% @return A tuple with the backend module and the backend state for this platform
-backend({StreamModule, Stream}) ->
+backend(StreamModule, Stream) ->
     BackendModule = ?MODULE:backend_module(),
     Variant = ?MODULE:variant(),
     BackendState = BackendModule:new(Variant, StreamModule, Stream),

libs/jit/src/jit_aarch64.erl

@@ -25,6 +25,7 @@
     new/3,
     stream/1,
     offset/1,
+    flush/1,
     debugger/1,
     used_regs/1,
     available_regs/1,
@@ -259,6 +260,16 @@ stream(#state{stream = Stream}) ->
 offset(#state{stream_module = StreamModule, stream = Stream}) ->
     StreamModule:offset(Stream).
 
+%%-----------------------------------------------------------------------------
+%% @doc Flush the current state (unused on aarch64)
+%% @end
+%% @param State current backend state
+%% @return The flushed state
+%%-----------------------------------------------------------------------------
+-spec flush(state()) -> state().
+flush(#state{} = State) ->
+    State.
+
 %%-----------------------------------------------------------------------------
 %% @doc Emit a debugger of breakpoint instruction. This is used for debugging
 %% and not in production.

libs/jit/src/jit_armv6m.erl

@@ -25,6 +25,7 @@
     new/3,
     stream/1,
     offset/1,
+    flush/1,
     debugger/1,
     used_regs/1,
     available_regs/1,
@@ -272,6 +273,16 @@ stream(#state{stream = Stream}) ->
 offset(#state{stream_module = StreamModule, stream = Stream}) ->
     StreamModule:offset(Stream).
 
+%%-----------------------------------------------------------------------------
+%% @doc Flush the current state, e.g. literal pools
+%% @end
+%% @param State current backend state
+%% @return The flushed state
+%%-----------------------------------------------------------------------------
+-spec flush(state()) -> state().
+flush(#state{} = State) ->
+    flush_literal_pool(State).
+
 %%-----------------------------------------------------------------------------
 %% @doc Emit a debugger of breakpoint instruction. This is used for debugging
 %% and not in production.

libs/jit/src/jit_stream_binary.erl

@@ -27,7 +27,8 @@
     offset/1,
     append/2,
     replace/3,
-    map/4
+    map/4,
+    flush/1
 ]).
 
 -export_type([stream/0]).
@@ -93,3 +94,14 @@ map(Stream, Offset, Length, MapFunction) ->
     {Prefix, <<Previous:Length/binary, Suffix/binary>>} = split_binary(Stream, Offset),
     Replacement = MapFunction(Previous),
     <<Prefix/binary, Replacement/binary, Suffix/binary>>.
+
+%%-----------------------------------------------------------------------------
+%% @param Stream        stream to flush
+%% @returns The stream flushed
+%% @doc     Flush the stream. NOP with binaries.
+%%
+%% @end
+%%-----------------------------------------------------------------------------
+-spec flush(stream()) -> stream().
+flush(Stream) ->
+    Stream.

libs/jit/src/jit_stream_mmap.erl

@@ -27,7 +27,8 @@
     offset/1,
     append/2,
     replace/3,
-    map/4
+    map/4,
+    flush/1
 ]).
 
 %% Additional nif
@@ -109,3 +110,14 @@ map(Stream, Offset, Length, MapFunction) ->
 -spec read(stream(), non_neg_integer(), pos_integer()) -> binary().
 read(_Stream, _Offset, _Length) ->
     erlang:nif_error(undefined).
+
+%%-----------------------------------------------------------------------------
+%% @param Stream        stream to flush
+%% @returns The stream flushed
+%% @doc     Flush the stream. Typically invalidates instruction cache.
+%%
+%% @end
+%%-----------------------------------------------------------------------------
+-spec flush(stream()) -> stream().
+flush(_Stream) ->
+    erlang:nif_error(undefined).

libs/jit/src/jit_x86_64.erl

@@ -25,6 +25,7 @@
     new/3,
     stream/1,
     offset/1,
+    flush/1,
     debugger/1,
     used_regs/1,
     available_regs/1,
@@ -244,6 +245,16 @@ stream(#state{stream = Stream}) ->
 offset(#state{stream_module = StreamModule, stream = Stream}) ->
     StreamModule:offset(Stream).
 
+%%-----------------------------------------------------------------------------
+%% @doc Flush the current state (unused on x86-64)
+%% @end
+%% @param State current backend state
+%% @return The flushed state
+%%-----------------------------------------------------------------------------
+-spec flush(state()) -> state().
+flush(#state{} = State) ->
+    State.
+
 %%-----------------------------------------------------------------------------
 %% @doc Emit a debugger of breakpoint instruction. This is used for debugging
 %% and not in production.

src/platforms/generic_unix/lib/jit_stream_mmap.c

@@ -193,6 +193,28 @@ static term nif_jit_stream_mmap_read(Context *ctx, int argc, term argv[])
     return term_from_literal_binary(js_obj->stream_base + offset, len, &ctx->heap, ctx->global);
 }
 
+static term nif_jit_stream_mmap_flush(Context *ctx, int argc, term argv[])
+{
+    UNUSED(argc);
+
+    void *js_obj_ptr;
+    if (UNLIKELY(!enif_get_resource(erl_nif_env_from_context(ctx), argv[0], jit_stream_mmap_resource_type, &js_obj_ptr))) {
+        RAISE_ERROR(BADARG_ATOM);
+    }
+    struct JITStreamMMap *js_obj = (struct JITStreamMMap *) js_obj_ptr;
+    if (IS_NULL_PTR(js_obj->stream_base)) {
+        RAISE_ERROR(BADARG_ATOM);
+    }
+
+#if defined(__APPLE__)
+    sys_icache_invalidate(js_obj->stream_base, js_obj->stream_size);
+#elif defined(__GNUC__)
+    __builtin___clear_cache(js_obj->stream_base, js_obj->stream_base + js_obj->stream_size);
+#endif
+
+    return argv[0];
+}
+
 static term nif_jit_stream_module(Context *ctx, int argc, term argv[])
 {
     UNUSED(argc);
@@ -226,6 +248,10 @@ static const struct Nif jit_stream_mmap_read_nif = {
     .base.type = NIFFunctionType,
     .nif_ptr = nif_jit_stream_mmap_read
 };
+static const struct Nif jit_stream_mmap_flush_nif = {
+    .base.type = NIFFunctionType,
+    .nif_ptr = nif_jit_stream_mmap_flush
+};
 
 ModuleNativeEntryPoint jit_stream_entry_point(Context *ctx, term jit_stream)
 {
@@ -239,18 +265,14 @@ ModuleNativeEntryPoint jit_stream_entry_point(Context *ctx, term jit_stream)
         return NULL;
     }
 
-#if defined(__APPLE__)
-    sys_icache_invalidate(js_obj->stream_base, js_obj->stream_size);
-#elif defined(__GNUC__)
-    __builtin___clear_cache(js_obj->stream_base, js_obj->stream_base + js_obj->stream_size);
-#endif
 #if JIT_ARCH_TARGET == JIT_ARCH_ARMV6M
     // Set thumb bit for armv6m
     ModuleNativeEntryPoint result = (ModuleNativeEntryPoint) js_obj->stream_base + 1;
 #else
     ModuleNativeEntryPoint result = (ModuleNativeEntryPoint) js_obj->stream_base;
 #endif
 
+    // Prevent module from being unmapped by dtor
     js_obj->stream_base = NULL;
     return result;
 }
@@ -291,6 +313,9 @@ const struct Nif *jit_stream_mmap_get_nif(const char *nifname)
         if (strcmp("read/3", rest) == 0) {
             return &jit_stream_mmap_read_nif;
         }
+        if (strcmp("flush/1", rest) == 0) {
+            return &jit_stream_mmap_flush_nif;
+        }
     }
     return NULL;
 }