JIT compile-time optimizations

Replace O(n) list lookups with O(log n) maps for hot data structures
in the JIT compiler, and batch consecutive small instruction emits
into single stream appends to reduce NIF dispatch overhead.

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

Author: pguyot

libs/jit/src/jit.erl

@@ -99,17 +99,22 @@
     tail_cache :: tail_cache()
 }).
 
--type tail_cache() :: [{tuple(), non_neg_integer()}] | disabled.
+-type tail_cache() :: #{tuple() => non_neg_integer()} | disabled.
 -type stream() :: any().
 
 %%-define(TRACE(Fmt, Args), io:format(Fmt, Args)).
 -define(TRACE(Fmt, Args), ok).
 
-tail_cache_find(_Key, disabled) -> false;
-tail_cache_find(Key, TC) -> lists:keyfind(Key, 1, TC).
+tail_cache_find(_Key, disabled) ->
+    false;
+tail_cache_find(Key, TC) ->
+    case TC of
+        #{Key := Value} -> {Key, Value};
+        _ -> false
+    end.
 
 tail_cache_store(_Key, _Value, disabled) -> disabled;
-tail_cache_store(Key, Value, TC) -> [{Key, Value} | TC].
+tail_cache_store(Key, Value, TC) -> TC#{Key => Value}.
 
 %%-define(ASSERT_ALL_NATIVE_FREE(St), MMod:assert_all_native_free(St)).
 %%-define(ASSERT(Expr), true = Expr).
@@ -154,11 +159,11 @@ compile(
             case erlang:function_exported(MMod, supports_tail_cache, 0) of
                 true ->
                     case MMod:supports_tail_cache() of
-                        true -> [];
+                        true -> #{};
                         false -> disabled
                     end;
                 false ->
-                    []
+                    #{}
             end
     },
     MSt1 = MMod:jump_table(MSt0, LabelsCount),

libs/jit/src/jit_aarch64.erl

@@ -154,11 +154,11 @@
     stream_module :: module(),
     stream :: stream(),
     offset :: non_neg_integer(),
-    branches :: [{non_neg_integer(), non_neg_integer(), non_neg_integer()}],
+    branches :: #{integer() | reference() => [{non_neg_integer(), non_neg_integer()}]},
     jump_table_start :: non_neg_integer(),
     available_regs :: non_neg_integer(),
     used_regs :: non_neg_integer(),
-    labels :: [{integer() | reference(), integer()}],
+    labels :: #{integer() | reference() => integer()},
     variant :: non_neg_integer(),
     regs :: jit_regs:regs()
 }).
@@ -286,12 +286,12 @@ new(Variant, StreamModule, Stream) ->
     #state{
         stream_module = StreamModule,
         stream = Stream,
-        branches = [],
+        branches = #{},
         jump_table_start = 0,
         offset = StreamModule:offset(Stream),
         available_regs = ?AVAILABLE_REGS_MASK,
         used_regs = 0,
-        labels = [],
+        labels = #{},
         variant = Variant,
         regs = jit_regs:new()
     }.
@@ -465,27 +465,25 @@ patch_branch(StreamModule, Stream, Offset, Type, LabelOffset) ->
 -spec patch_branches_for_label(
     module(),
     stream(),
-    integer(),
+    integer() | reference(),
     non_neg_integer(),
-    [{integer(), non_neg_integer(), any()}]
-) -> {stream(), [{integer(), non_neg_integer(), any()}]}.
+    #{integer() | reference() => [{non_neg_integer(), non_neg_integer()}]}
+) ->
+    {stream(), #{integer() | reference() => [{non_neg_integer(), non_neg_integer()}]}}.
 patch_branches_for_label(StreamModule, Stream, TargetLabel, LabelOffset, Branches) ->
-    patch_branches_for_label(StreamModule, Stream, TargetLabel, LabelOffset, Branches, []).
-
-patch_branches_for_label(_StreamModule, Stream, _TargetLabel, _LabelOffset, [], Acc) ->
-    {Stream, lists:reverse(Acc)};
-patch_branches_for_label(
-    StreamModule,
-    Stream0,
-    TargetLabel,
-    LabelOffset,
-    [{Label, Offset, Type} | Rest],
-    Acc
-) when Label =:= TargetLabel ->
-    Stream1 = patch_branch(StreamModule, Stream0, Offset, Type, LabelOffset),
-    patch_branches_for_label(StreamModule, Stream1, TargetLabel, LabelOffset, Rest, Acc);
-patch_branches_for_label(StreamModule, Stream, TargetLabel, LabelOffset, [Branch | Rest], Acc) ->
-    patch_branches_for_label(StreamModule, Stream, TargetLabel, LabelOffset, Rest, [Branch | Acc]).
+    case Branches of
+        #{TargetLabel := BrList} ->
+            Stream1 = lists:foldl(
+                fun({Offset, Type}, AccStream) ->
+                    patch_branch(StreamModule, AccStream, Offset, Type, LabelOffset)
+                end,
+                Stream,
+                BrList
+            ),
+            {Stream1, maps:remove(TargetLabel, Branches)};
+        _ ->
+            {Stream, Branches}
+    end.
 
 %%-----------------------------------------------------------------------------
 %% @doc Rewrite stream to update all branches for labels.
@@ -494,19 +492,29 @@ patch_branches_for_label(StreamModule, Stream, TargetLabel, LabelOffset, [Branch
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
 -spec update_branches(state()) -> state().
-update_branches(#state{branches = []} = State) ->
-    State;
 update_branches(
     #state{
         stream_module = StreamModule,
         stream = Stream0,
-        branches = [{Label, Offset, Type} | BranchesT],
+        branches = Branches,
         labels = Labels
     } = State
 ) ->
-    {Label, LabelOffset} = lists:keyfind(Label, 1, Labels),
-    Stream1 = patch_branch(StreamModule, Stream0, Offset, Type, LabelOffset),
-    update_branches(State#state{stream = Stream1, branches = BranchesT}).
+    Stream1 = maps:fold(
+        fun(Label, BrList, AccStream) ->
+            #{Label := LabelOffset} = Labels,
+            lists:foldl(
+                fun({Offset, Type}, AccStream2) ->
+                    patch_branch(StreamModule, AccStream2, Offset, Type, LabelOffset)
+                end,
+                AccStream,
+                BrList
+            )
+        end,
+        Stream0,
+        Branches
+    ),
+    State#state{stream = Stream1, branches = #{}}.
 
 %%-----------------------------------------------------------------------------
 %% @doc Emit a call (call with return) to a primitive with arguments. This
@@ -650,21 +658,22 @@ jump_to_label(
     Label
 ) ->
     Offset = StreamModule:offset(Stream0),
-    case lists:keyfind(Label, 1, Labels) of
-        {Label, LabelOffset} ->
+    case Labels of
+        #{Label := LabelOffset} ->
             % Label is already known, emit direct branch without relocation
             Rel = LabelOffset - Offset,
             I1 = jit_aarch64_asm:b(Rel),
             Stream1 = StreamModule:append(Stream0, I1),
             State#state{stream = Stream1, regs = jit_regs:unreachable(State#state.regs)};
-        false ->
+        _ ->
             % Label not yet known, emit placeholder and add relocation
             I1 = jit_aarch64_asm:b(0),
-            Reloc = {Label, Offset, b},
+            BrEntry = {Offset, b},
+            ExistingBrs = maps:get(Label, AccBranches, []),
             Stream1 = StreamModule:append(Stream0, I1),
             State#state{
                 stream = Stream1,
-                branches = [Reloc | AccBranches],
+                branches = AccBranches#{Label => [BrEntry | ExistingBrs]},
                 regs = jit_regs:unreachable(State#state.regs)
             }
     end.
@@ -2354,21 +2363,26 @@ set_continuation_to_label(
     Temp = first_avail(Avail),
     Offset = StreamModule:offset(Stream0),
     Regs1 = jit_regs:invalidate_reg(Regs0, Temp),
-    case lists:keyfind(Label, 1, Labels) of
-        {Label, LabelOffset} ->
+    case Labels of
+        #{Label := LabelOffset} ->
             Rel = LabelOffset - Offset,
             I1 = jit_aarch64_asm:adr(Temp, Rel),
             I2 = jit_aarch64_asm:str(Temp, ?JITSTATE_CONTINUATION),
             Code = <<I1/binary, I2/binary>>,
             Stream1 = StreamModule:append(Stream0, Code),
             State#state{stream = Stream1, regs = Regs1};
-        false ->
+        _ ->
             I1 = jit_aarch64_asm:adr(Temp, 0),
-            Reloc = {Label, Offset, {adr, Temp}},
+            BrEntry = {Offset, {adr, Temp}},
             I2 = jit_aarch64_asm:str(Temp, ?JITSTATE_CONTINUATION),
             Code = <<I1/binary, I2/binary>>,
             Stream1 = StreamModule:append(Stream0, Code),
-            State#state{stream = Stream1, branches = [Reloc | Branches], regs = Regs1}
+            ExistingBrs = maps:get(Label, Branches, []),
+            State#state{
+                stream = Stream1,
+                branches = Branches#{Label => [BrEntry | ExistingBrs]},
+                regs = Regs1
+            }
     end.
 
 %%-----------------------------------------------------------------------------
@@ -2393,12 +2407,19 @@ set_continuation_to_offset(
     OffsetRef = make_ref(),
     Offset = StreamModule:offset(Stream0),
     I1 = jit_aarch64_asm:adr(Temp, 0),
-    Reloc = {OffsetRef, Offset, {adr, Temp}},
+    BrEntry = {Offset, {adr, Temp}},
     I2 = jit_aarch64_asm:str(Temp, ?JITSTATE_CONTINUATION),
     Code = <<I1/binary, I2/binary>>,
     Stream1 = StreamModule:append(Stream0, Code),
     Regs1 = jit_regs:invalidate_reg(Regs0, Temp),
-    {State#state{stream = Stream1, branches = [Reloc | Branches], regs = Regs1}, OffsetRef}.
+    {
+        State#state{
+            stream = Stream1,
+            branches = Branches#{OffsetRef => [BrEntry]},
+            regs = Regs1
+        },
+        OffsetRef
+    }.
 
 %%-----------------------------------------------------------------------------
 %% @doc Implement a continuation entry point. On AArch64 this is a nop
@@ -2783,20 +2804,24 @@ call_only_or_schedule_next(
     Stream1 = StreamModule:append(Stream0, <<I1/binary, I2/binary, I3/binary>>),
     BNEOffset = StreamModule:offset(Stream1),
 
-    case lists:keyfind(Label, 1, Labels) of
-        {Label, LabelOffset} ->
+    case Labels of
+        #{Label := LabelOffset} ->
             % Label is already known, emit direct branch with calculated offset
             % Calculate relative offset (must be 4-byte aligned)
             Rel = LabelOffset - BNEOffset,
             I4 = jit_aarch64_asm:bcc(ne, Rel),
             Stream2 = StreamModule:append(Stream1, I4),
             State1 = State0#state{stream = Stream2};
-        false ->
+        _ ->
             % Label not yet known, emit placeholder and add relocation
             I4 = jit_aarch64_asm:bcc(ne, 0),
-            Reloc1 = {Label, BNEOffset, {bcc, ne}},
+            BrEntry = {BNEOffset, {bcc, ne}},
+            ExistingBrs = maps:get(Label, Branches, []),
             Stream2 = StreamModule:append(Stream1, I4),
-            State1 = State0#state{stream = Stream2, branches = [Reloc1 | Branches]}
+            State1 = State0#state{
+                stream = Stream2,
+                branches = Branches#{Label => [BrEntry | ExistingBrs]}
+            }
     end,
     State2 = set_continuation_to_label(State1, Label),
     call_primitive_last(State2, ?PRIM_SCHEDULE_NEXT_CP, [ctx, jit_state]).
@@ -2894,7 +2919,7 @@ return_labels_and_lines(
 ) ->
     SortedLabels = lists:keysort(2, [
         {Label, LabelOffset}
-     || {Label, LabelOffset} <- Labels, is_integer(Label)
+     || {Label, LabelOffset} <- maps:to_list(Labels), is_integer(Label)
     ]),
 
     I1 = jit_aarch64_asm:adr(r0, 8),
@@ -3049,10 +3074,10 @@ add_label(
     ),
 
     State#state{
-        stream = Stream2, branches = RemainingBranches, labels = [{Label, LabelOffset} | Labels]
+        stream = Stream2, branches = RemainingBranches, labels = Labels#{Label => LabelOffset}
     };
 add_label(#state{labels = Labels} = State, Label, Offset) ->
-    State#state{labels = [{Label, Offset} | Labels]}.
+    State#state{labels = Labels#{Label => Offset}}.
 
 -ifdef(JIT_DWARF).
 %%-----------------------------------------------------------------------------