JIT: Optimize memory usage by rewriting jump table immediately

libs/estdlib/src/code_server.erl

@@ -152,7 +152,7 @@ set_native_code(_Module, _LabelsCount, _Stream) ->
 load(Module) ->
     case erlang:system_info(emu_flavor) of
         jit ->
-            % atomvm_heap_growth, fibonacci divides compilation time by two
+            % atomvm_heap_growth, fibonacci reduces compilation time
             {Pid, Ref} = spawn_opt(
                 fun() ->
                     try

libs/jit/src/jit_aarch64.erl

@@ -134,6 +134,7 @@
     stream :: stream(),
     offset :: non_neg_integer(),
     branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
+    jump_table_start :: non_neg_integer(),
     available_regs :: [aarch64_register()],
     used_regs :: [aarch64_register()],
     labels :: [{integer() | reference(), integer()}],
@@ -233,6 +234,7 @@ new(Variant, StreamModule, Stream) ->
         stream_module = StreamModule,
         stream = Stream,
         branches = [],
+        jump_table_start = 0,
         offset = StreamModule:offset(Stream),
         available_regs = ?AVAILABLE_REGS,
         used_regs = [],
@@ -355,22 +357,21 @@ assert_all_native_free(#state{
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec jump_table(state(), pos_integer()) -> state().
-jump_table(State, LabelsCount) ->
-    jump_table0(State, 0, LabelsCount).
+jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
+    JumpTableStart = StreamModule:offset(Stream0),
+    jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
 
 -spec jump_table0(state(), non_neg_integer(), pos_integer()) -> state().
 jump_table0(State, N, LabelsCount) when N > LabelsCount ->
     State;
 jump_table0(
-    #state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
+    #state{stream_module = StreamModule, stream = Stream0} = State,
     N,
     LabelsCount
 ) ->
-    Offset = StreamModule:offset(Stream0),
     BranchInstr = jit_aarch64_asm:b(0),
-    Reloc = {N, Offset, b},
     Stream1 = StreamModule:append(Stream0, BranchInstr),
-    jump_table0(State#state{stream = Stream1, branches = [Reloc | Branches]}, N + 1, LabelsCount).
+    jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
 
 %%-----------------------------------------------------------------------------
 %% @doc Rewrite stream to update all branches for labels.
@@ -2343,5 +2344,22 @@ add_label(#state{stream_module = StreamModule, stream = Stream} = State, Label)
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec add_label(state(), integer() | reference(), integer()) -> state().
+add_label(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        jump_table_start = JumpTableStart,
+        labels = Labels
+    } = State,
+    Label,
+    LabelOffset
+) when is_integer(Label) ->
+    % Patch the jump table entry immediately
+    % Each b instruction is 4 bytes
+    JumpTableEntryOffset = JumpTableStart + Label * 4,
+    RelativeOffset = LabelOffset - JumpTableEntryOffset,
+    BranchInstr = jit_aarch64_asm:b(RelativeOffset),
+    Stream1 = StreamModule:replace(Stream0, JumpTableEntryOffset, BranchInstr),
+    State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
 add_label(#state{labels = Labels} = State, Label, Offset) ->
     State#state{labels = [{Label, Offset} | Labels]}.

libs/jit/src/jit_armv6m.erl

@@ -134,6 +134,7 @@
     stream :: stream(),
     offset :: non_neg_integer(),
     branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
+    jump_table_start :: non_neg_integer(),
     available_regs :: [armv6m_register()],
     used_regs :: [armv6m_register()],
     labels :: [{integer() | reference(), integer()}],
@@ -247,6 +248,7 @@ new(Variant, StreamModule, Stream) ->
         stream_module = StreamModule,
         stream = Stream,
         branches = [],
+        jump_table_start = 0,
         offset = StreamModule:offset(Stream),
         available_regs = ?AVAILABLE_REGS,
         used_regs = [],
@@ -380,13 +382,14 @@ assert_all_native_free(#state{
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec jump_table(state(), pos_integer()) -> state().
-jump_table(State, LabelsCount) ->
-    jump_table0(State, 0, LabelsCount).
+jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
+    JumpTableStart = StreamModule:offset(Stream0),
+    jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
 
 jump_table0(State, N, LabelsCount) when N > LabelsCount ->
     State;
 jump_table0(
-    #state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
+    #state{stream_module = StreamModule, stream = Stream0} = State,
     N,
     LabelsCount
 ) ->
@@ -399,15 +402,7 @@ jump_table0(
     JumpEntry = <<I1/binary, I2/binary, I3/binary, I4/binary, 16#FFFFFFFF:32>>,
     Stream1 = StreamModule:append(Stream0, JumpEntry),
 
-    % Add relocation for the data entry so update_branches/2 can patch the jump target
-    DataOffset = StreamModule:offset(Stream1) - 4,
-    % Calculate the offset of the add instruction (3rd instruction, at offset 4 from entry start)
-    EntryStartOffset = StreamModule:offset(Stream1) - 12,
-    AddInstrOffset = EntryStartOffset + 4,
-    DataReloc = {N, DataOffset, {jump_table_data, AddInstrOffset}},
-    UpdatedState = State#state{stream = Stream1, branches = [DataReloc | Branches]},
-
-    jump_table0(UpdatedState, N + 1, LabelsCount).
+    jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
 
 %%-----------------------------------------------------------------------------
 %% @doc Rewrite stream to update all branches for labels.
@@ -500,13 +495,7 @@ update_branches(
                                 I4 = <<RelativeOffset:32/little>>,
                                 <<I1/binary, I2/binary, I3/binary, I4/binary>>
                         end
-                end;
-            {jump_table_data, AddInstrOffset} ->
-                % Calculate offset from 'add pc, pc, r3' instruction + 4 to target label
-                % PC when add instruction executes
-                AddPC = AddInstrOffset + 4,
-                RelativeOffset = LabelOffset - AddPC,
-                <<RelativeOffset:32/little>>
+                end
         end,
     Stream1 = StreamModule:replace(Stream0, Offset, NewInstr),
     update_branches(State#state{stream = Stream1, branches = BranchesT}).
@@ -3288,5 +3277,34 @@ add_label(#state{stream_module = StreamModule, stream = Stream0} = State0, Label
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec add_label(state(), integer() | reference(), integer()) -> state().
+add_label(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        jump_table_start = JumpTableStart,
+        labels = Labels
+    } = State,
+    Label,
+    LabelOffset
+) when is_integer(Label) ->
+    % Patch the jump table entry immediately
+    % Each jump table entry is 12 bytes:
+    % - ldr r3, [pc, 4] (2 bytes) at offset 0
+    % - push {...} (2 bytes) at offset 2
+    % - add pc, r3 (2 bytes) at offset 4
+    % - nop (2 bytes) at offset 6
+    % - data (4 bytes) at offset 8
+    JumpTableEntryStart = JumpTableStart + Label * 12,
+    DataOffset = JumpTableEntryStart + 8,
+    AddInstrOffset = JumpTableEntryStart + 4,
+
+    % Calculate offset from 'add pc, pc, r3' instruction + 4 to target label
+    % PC when add instruction executes
+    AddPC = AddInstrOffset + 4,
+    RelativeOffset = LabelOffset - AddPC,
+    DataBytes = <<RelativeOffset:32/little>>,
+
+    Stream1 = StreamModule:replace(Stream0, DataOffset, DataBytes),
+    State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
 add_label(#state{labels = Labels} = State, Label, Offset) ->
     State#state{labels = [{Label, Offset} | Labels]}.

libs/jit/src/jit_x86_64.erl

@@ -115,6 +115,7 @@
     stream :: stream(),
     offset :: non_neg_integer(),
     branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
+    jump_table_start :: non_neg_integer(),
     available_regs :: [x86_64_register()],
     used_regs :: [x86_64_register()],
     labels :: [{integer() | reference(), integer()}],
@@ -218,6 +219,7 @@ new(Variant, StreamModule, Stream) ->
         stream_module = StreamModule,
         stream = Stream,
         branches = [],
+        jump_table_start = 0,
         offset = StreamModule:offset(Stream),
         available_regs = ?AVAILABLE_REGS,
         used_regs = [],
@@ -340,21 +342,21 @@ assert_all_native_free(State) ->
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec jump_table(state(), pos_integer()) -> state().
-jump_table(State, LabelsCount) ->
-    jump_table0(State, 0, LabelsCount).
+jump_table(#state{stream_module = StreamModule, stream = Stream0} = State, LabelsCount) ->
+    JumpTableStart = StreamModule:offset(Stream0),
+    jump_table0(State#state{jump_table_start = JumpTableStart}, 0, LabelsCount).
 
 jump_table0(State, N, LabelsCount) when N > LabelsCount ->
     State;
 jump_table0(
-    #state{stream_module = StreamModule, stream = Stream0, branches = Branches} = State,
+    #state{stream_module = StreamModule, stream = Stream0} = State,
     N,
     LabelsCount
 ) ->
-    Offset = StreamModule:offset(Stream0),
-    {RelocOffset, I1} = jit_x86_64_asm:jmp_rel32(1),
-    Reloc = {N, Offset + RelocOffset, 32},
+    % Placeholder, encodes with 0xffffffff
+    {_RelocOffset, I1} = jit_x86_64_asm:jmp_rel32(4),
     Stream1 = StreamModule:append(Stream0, I1),
-    jump_table0(State#state{stream = Stream1, branches = [Reloc | Branches]}, N + 1, LabelsCount).
+    jump_table0(State#state{stream = Stream1}, N + 1, LabelsCount).
 
 %%-----------------------------------------------------------------------------
 %% @doc Rewrite stream to update all branches for labels.
@@ -2086,5 +2088,22 @@ add_label(#state{stream_module = StreamModule, stream = Stream} = State, Label)
     add_label(State, Label, Offset).
 
 -spec add_label(state(), integer() | reference(), integer()) -> state().
+add_label(
+    #state{
+        stream_module = StreamModule,
+        stream = Stream0,
+        jump_table_start = JumpTableStart,
+        labels = Labels
+    } = State,
+    Label,
+    LabelOffset
+) when is_integer(Label) ->
+    % Patch the jump table entry immediately
+    % Each jmp_rel32 instruction is 5 bytes
+    JumpTableEntryOffset = JumpTableStart + Label * 5,
+    RelativeOffset = LabelOffset - JumpTableEntryOffset,
+    {_RelocOffset, JmpInstruction} = jit_x86_64_asm:jmp_rel32(RelativeOffset),
+    Stream1 = StreamModule:replace(Stream0, JumpTableEntryOffset, JmpInstruction),
+    State#state{stream = Stream1, labels = [{Label, LabelOffset} | Labels]};
 add_label(#state{labels = Labels} = State, Label, Offset) ->
     State#state{labels = [{Label, Offset} | Labels]}.