JIT: factorize tail calls to reduce binary size

Use a cache to remember tail calls that were already implemented and
replace further implementations of the same tail call with a jump to the
previous implementation.

Coverage shows that all cases are covered in libs/estdlib/src and libs/jit/src:

OP_RETURN: 50 misses, 1735 hits (97%)
OP_JUMP/OP_CALL_LAST/OP_CALL_ONLY: 656 misses, 389 hits (37%)
OP_CALL_LAST: 220 misses, 206 hits (48%)
OP_FUNC_INFO: 58 misses, 1619 hits (97%)

Signed-off-by: Paul Guyot <[email protected]>

Author: pguyot

libs/jit/src/jit.erl

@@ -100,7 +100,8 @@
     labels_count :: pos_integer(),
     atom_resolver :: fun((integer()) -> atom()),
     literal_resolver :: fun((integer()) -> any()),
-    type_resolver :: fun((integer()) -> any())
+    type_resolver :: fun((integer()) -> any()),
+    tail_cache :: [{tuple(), non_neg_integer()}]
 }).
 
 -type stream() :: any().
@@ -142,7 +143,8 @@ compile(
         labels_count = LabelsCount,
         atom_resolver = AtomResolver,
         literal_resolver = LiteralResolver,
-        type_resolver = TypeResolver
+        type_resolver = TypeResolver,
+        tail_cache = []
     },
     {State1, MSt2} = first_pass(Opcodes, MMod, MSt1, State0),
     MSt3 = second_pass(MMod, MSt2, State1),
@@ -170,18 +172,30 @@ first_pass(
     ?ASSERT_ALL_NATIVE_FREE(MSt1),
     first_pass(Rest1, MMod, MSt1, State0);
 % 2
-first_pass(<<?OP_FUNC_INFO, Rest0/binary>>, MMod, MSt0, State0) ->
+first_pass(<<?OP_FUNC_INFO, Rest0/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {_ModuleAtom, Rest1} = decode_atom(Rest0),
     {_FunctionName, Rest2} = decode_atom(Rest1),
     {_Arity, Rest3} = decode_literal(Rest2),
     ?TRACE("OP_FUNC_INFO ~p, ~p, ~p\n", [_ModuleAtom, _FunctionName, _Arity]),
-    % Implement function clause at the previous label. (TODO: optimize it out to save space)
-    MSt1 = MMod:call_primitive_last(MSt0, ?PRIM_RAISE_ERROR, [
-        ctx, jit_state, offset, ?FUNCTION_CLAUSE_ATOM
-    ]),
-    ?ASSERT_ALL_NATIVE_FREE(MSt1),
-    first_pass(Rest3, MMod, MSt1, State0);
+    % Implement function clause at the previous label.
+    Offset = MMod:offset(MSt0),
+    {MSt1, OffsetReg} = MMod:move_to_native_register(MSt0, Offset),
+    TailCacheKey = {call_primitive_last, ?PRIM_RAISE_ERROR, [OffsetReg, ?FUNCTION_CLAUSE_ATOM]},
+    State1 =
+        case lists:keyfind(TailCacheKey, 1, TC) of
+            false ->
+                MSt3 = MMod:call_primitive_last(MSt1, ?PRIM_RAISE_ERROR, [
+                    ctx, jit_state, {free, OffsetReg}, ?FUNCTION_CLAUSE_ATOM
+                ]),
+                State0#state{tail_cache = [{TailCacheKey, Offset} | TC]};
+            {TailCacheKey, CacheOffset} ->
+                MSt2 = MMod:jump_to_offset(MSt1, CacheOffset),
+                MSt3 = MMod:free_native_registers(MSt2, [OffsetReg]),
+                State0
+        end,
+    ?ASSERT_ALL_NATIVE_FREE(MSt3),
+    first_pass(Rest3, MMod, MSt3, State1);
 % 3
 first_pass(
     <<?OP_INT_CALL_END>>, MMod, MSt0, #state{labels_count = LabelsCount} = State
@@ -203,26 +217,56 @@ first_pass(<<?OP_CALL, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt1),
     first_pass(Rest2, MMod, MSt1, State0);
 % 5
-first_pass(<<?OP_CALL_LAST, Rest0/binary>>, MMod, MSt0, State0) ->
+first_pass(<<?OP_CALL_LAST, Rest0/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {_Arity, Rest1} = decode_literal(Rest0),
     {Label, Rest2} = decode_label(Rest1),
     {NWords, Rest3} = decode_literal(Rest2),
     ?TRACE("OP_CALL_LAST ~p, ~p, ~p\n", [_Arity, Label, NWords]),
-    MSt1 = MMod:move_to_cp(MSt0, {y_reg, NWords}),
-    MSt2 = MMod:increment_sp(MSt1, NWords + 1),
-    MSt3 = MMod:call_only_or_schedule_next(MSt2, Label),
+    TailCacheKey0 = {op_call_last, NWords, Label},
+    case lists:keyfind(TailCacheKey0, 1, TC) of
+        false ->
+            Offset0 = MMod:offset(MSt0),
+            MSt1 = MMod:move_to_cp(MSt0, {y_reg, NWords}),
+            MSt2 = MMod:increment_sp(MSt1, NWords + 1),
+            TailCacheKey1 = {op_call_only, Label},
+            case lists:keyfind(TailCacheKey1, 1, TC) of
+                false ->
+                    Offset1 = MMod:offset(MSt2),
+                    MSt3 = MMod:call_only_or_schedule_next(MSt2, Label),
+                    State1 = State0#state{
+                        tail_cache = [{TailCacheKey1, Offset1}, {TailCacheKey0, Offset0} | TC]
+                    };
+                {TailCacheKey1, Offset1} ->
+                    MSt3 = MMod:jump_to_offset(MSt2, Offset1),
+                    State1 = State0#state{
+                        tail_cache = [{TailCacheKey0, Offset0} | TC]
+                    }
+            end;
+        {TailCacheKey0, Offset0} ->
+            MSt3 = MMod:jump_to_offset(MSt0, Offset0),
+            State1 = State0
+    end,
     ?ASSERT_ALL_NATIVE_FREE(MSt3),
-    first_pass(Rest3, MMod, MSt3, State0);
+    first_pass(Rest3, MMod, MSt3, State1);
 % 6
-first_pass(<<?OP_CALL_ONLY, Rest0/binary>>, MMod, MSt0, State0) ->
+first_pass(<<?OP_CALL_ONLY, Rest0/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {_Arity, Rest1} = decode_literal(Rest0),
     {Label, Rest2} = decode_label(Rest1),
     ?TRACE("OP_CALL_ONLY ~p, ~p\n", [_Arity, Label]),
-    MSt1 = MMod:call_only_or_schedule_next(MSt0, Label),
+    TailCacheKey = {op_call_only, Label},
+    case lists:keyfind(TailCacheKey, 1, TC) of
+        false ->
+            Offset = MMod:offset(MSt0),
+            MSt1 = MMod:call_only_or_schedule_next(MSt0, Label),
+            State1 = State0#state{tail_cache = [{TailCacheKey, Offset} | TC]};
+        {TailCacheKey, Offset} ->
+            MSt1 = MMod:jump_to_offset(MSt0, Offset),
+            State1 = State0
+    end,
     ?ASSERT_ALL_NATIVE_FREE(MSt1),
-    first_pass(Rest2, MMod, MSt1, State0);
+    first_pass(Rest2, MMod, MSt1, State1);
 % 7
 first_pass(<<?OP_CALL_EXT, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -348,7 +392,7 @@ first_pass(<<?OP_DEALLOCATE, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt2),
     first_pass(Rest1, MMod, MSt2, State0);
 % 19
-first_pass(<<?OP_RETURN, Rest/binary>>, MMod, MSt0, State0) ->
+first_pass(<<?OP_RETURN, Rest/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     ?TRACE("OP_RETURN\n", []),
     % Optimized return: check if returning within same module
@@ -371,9 +415,18 @@ first_pass(<<?OP_RETURN, Rest/binary>>, MMod, MSt0, State0) ->
     ),
     MSt5 = MMod:free_native_registers(MSt4, [CpReg0]),
     % Different module: use existing slow path
-    MSt6 = MMod:call_primitive_last(MSt5, ?PRIM_RETURN, [ctx, jit_state]),
+    TailCacheKey = {call_primitive_last, ?PRIM_RETURN},
+    case lists:keyfind(TailCacheKey, 1, TC) of
+        false ->
+            Offset = MMod:offset(MSt5),
+            MSt6 = MMod:call_primitive_last(MSt5, ?PRIM_RETURN, [ctx, jit_state]),
+            State1 = State0#state{tail_cache = [{TailCacheKey, Offset} | TC]};
+        {TailCacheKey, Offset} ->
+            MSt6 = MMod:jump_to_offset(MSt5, Offset),
+            State1 = State0
+    end,
     ?ASSERT_ALL_NATIVE_FREE(MSt6),
-    first_pass(Rest, MMod, MSt6, State0);
+    first_pass(Rest, MMod, MSt6, State1);
 % 20
 first_pass(<<?OP_SEND, Rest/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
@@ -836,13 +889,22 @@ first_pass(<<?OP_SELECT_TUPLE_ARITY, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt5),
     first_pass(Rest4, MMod, MSt5, State0);
 % 61
-first_pass(<<?OP_JUMP, Rest0/binary>>, MMod, MSt0, State0) ->
+first_pass(<<?OP_JUMP, Rest0/binary>>, MMod, MSt0, #state{tail_cache = TC} = State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {Label, Rest1} = decode_label(Rest0),
     ?TRACE("OP_JUMP ~p\n", [Label]),
-    MSt1 = MMod:call_only_or_schedule_next(MSt0, Label),
-    ?ASSERT_ALL_NATIVE_FREE(MSt1),
-    first_pass(Rest1, MMod, MSt1, State0);
+    TailCacheKey = {op_call_only, Label},
+    case lists:keyfind(TailCacheKey, 1, TC) of
+        false ->
+            Offset = MMod:offset(MSt0),
+            MSt1 = MMod:call_only_or_schedule_next(MSt0, Label),
+            ?ASSERT_ALL_NATIVE_FREE(MSt1),
+            first_pass(Rest1, MMod, MSt1, State0#state{tail_cache = [{TailCacheKey, Offset} | TC]});
+        {TailCacheKey, Offset} ->
+            MSt1 = MMod:jump_to_offset(MSt0, Offset),
+            ?ASSERT_ALL_NATIVE_FREE(MSt1),
+            first_pass(Rest1, MMod, MSt1, State0)
+    end;
 % 62
 % Same implementation as OP_TRY, to confirm.
 first_pass(<<?OP_CATCH, Rest0/binary>>, MMod, MSt0, State0) ->

libs/jit/src/jit_aarch64.erl

@@ -38,6 +38,7 @@
     return_if_not_equal_to_ctx/2,
     jump_to_label/2,
     jump_to_continuation/2,
+    jump_to_offset/2,
     if_block/3,
     if_else_block/4,
     shift_right/3,
@@ -531,6 +532,13 @@ jump_to_label(
             State#state{stream = Stream1, branches = [Reloc | AccBranches]}
     end.
 
+jump_to_offset(#state{stream_module = StreamModule, stream = Stream0} = State, TargetOffset) ->
+    Offset = StreamModule:offset(Stream0),
+    Rel = TargetOffset - Offset,
+    I1 = jit_aarch64_asm:b(Rel),
+    Stream1 = StreamModule:append(Stream0, I1),
+    State#state{stream = Stream1}.
+
 %%-----------------------------------------------------------------------------
 %% @doc Jump to a continuation address stored in a register.
 %% This is used for optimized intra-module returns.

libs/jit/src/jit_armv6m.erl

@@ -38,6 +38,7 @@
     return_if_not_equal_to_ctx/2,
     jump_to_label/2,
     jump_to_continuation/2,
+    jump_to_offset/2,
     if_block/3,
     if_else_block/4,
     shift_right/3,
@@ -731,6 +732,12 @@ jump_to_label(
     Stream1 = StreamModule:append(Stream0, CodeBlock),
     State1#state{stream = Stream1}.
 
+jump_to_offset(#state{stream_module = StreamModule, stream = Stream0} = State, TargetOffset) ->
+    Offset = StreamModule:offset(Stream0),
+    CodeBlock = branch_to_offset_code(State, Offset, TargetOffset),
+    Stream1 = StreamModule:append(Stream0, CodeBlock),
+    State#state{stream = Stream1}.
+
 %%-----------------------------------------------------------------------------
 %% @doc Jump to address in continuation pointer register
 %% The continuation points to a function prologue, so we need to compute
@@ -793,15 +800,14 @@ jump_to_continuation(
     % Free all registers as this is a terminal instruction
     State1#state{stream = Stream2, available_regs = ?AVAILABLE_REGS, used_regs = []}.
 
-branch_to_label_code(State, Offset, Label, {Label, LabelOffset}) when
-    LabelOffset - Offset =< 2050, LabelOffset - Offset >= -2044
+branch_to_offset_code(_State, Offset, TargetOffset) when
+    TargetOffset - Offset =< 2050, TargetOffset - Offset >= -2044
 ->
     % Near branch: use direct B instruction
-    Rel = LabelOffset - Offset,
-    CodeBlock = jit_armv6m_asm:b(Rel),
-    {State, CodeBlock};
-branch_to_label_code(
-    #state{available_regs = [TempReg | _]} = State0, Offset, Label, {Label, LabelOffset}
+    Rel = TargetOffset - Offset,
+    jit_armv6m_asm:b(Rel);
+branch_to_offset_code(
+    #state{available_regs = [TempReg | _]}, Offset, TargetOffset
 ) ->
     % Far branch: use register-based sequence, need temporary register
     if
@@ -812,19 +818,22 @@ branch_to_label_code(
             I3 = jit_armv6m_asm:bx(TempReg),
             % Unaligned : need nop
             I4 = jit_armv6m_asm:nop(),
-            LiteralValue = LabelOffset - Offset - 5,
+            LiteralValue = TargetOffset - Offset - 5,
             I5 = <<LiteralValue:32/little>>,
-            CodeBlock = <<I1/binary, I2/binary, I3/binary, I4/binary, I5/binary>>;
+            <<I1/binary, I2/binary, I3/binary, I4/binary, I5/binary>>;
         true ->
             % Unaligned
             I1 = jit_armv6m_asm:ldr(TempReg, {pc, 4}),
             I2 = jit_armv6m_asm:add(TempReg, pc),
             I3 = jit_armv6m_asm:bx(TempReg),
-            LiteralValue = LabelOffset - Offset - 5,
+            LiteralValue = TargetOffset - Offset - 5,
             I4 = <<LiteralValue:32/little>>,
-            CodeBlock = <<I1/binary, I2/binary, I3/binary, I4/binary>>
-    end,
-    {State0, CodeBlock};
+            <<I1/binary, I2/binary, I3/binary, I4/binary>>
+    end.
+
+branch_to_label_code(State, Offset, Label, {Label, LabelOffset}) ->
+    CodeBlock = branch_to_offset_code(State, Offset, LabelOffset),
+    {State, CodeBlock};
 branch_to_label_code(
     #state{available_regs = [TempReg | _], branches = Branches} = State0, Offset, Label, false
 ) ->

libs/jit/src/jit_x86_64.erl

@@ -38,6 +38,7 @@
     return_if_not_equal_to_ctx/2,
     jump_to_label/2,
     jump_to_continuation/2,
+    jump_to_offset/2,
     if_block/3,
     if_else_block/4,
     shift_right/3,
@@ -524,6 +525,13 @@ jump_to_label(
             State#state{stream = Stream1, branches = [Reloc | AccBranches]}
     end.
 
+jump_to_offset(#state{stream_module = StreamModule, stream = Stream0} = State, TargetOffset) ->
+    Offset = StreamModule:offset(Stream0),
+    RelOffset = TargetOffset - Offset,
+    I1 = jit_x86_64_asm:jmp(RelOffset),
+    Stream1 = StreamModule:append(Stream0, I1),
+    State#state{stream = Stream1}.
+
 %%-----------------------------------------------------------------------------
 %% @doc Jump to a continuation address stored in a register.
 %% This is used for optimized intra-module returns.