Merge pull request #1834 from pguyot/w39/jit-optimize-term_to_int-using-types

JIT: optimize term to int using types

These changes are made under both the "Apache 2.0" and the "GNU Lesser General
Public License 2.1 or later" license terms (dual license).

SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later

Author: bettio

libs/estdlib/src/code_server.erl

@@ -38,6 +38,7 @@
     code_chunk/1,
     atom_resolver/2,
     literal_resolver/2,
+    type_resolver/2,
     set_native_code/3
 ]).
 
@@ -126,6 +127,14 @@ atom_resolver(_Module, _Index) ->
 literal_resolver(_Module, _Index) ->
     erlang:nif_error(undefined).
 
+%% @doc Get a type from its index
+%% @return The type information
+%% @param Module module get a type from
+%% @param Index type index in the module
+-spec type_resolver(Module :: module(), Index :: non_neg_integer()) -> any().
+type_resolver(_Module, _Index) ->
+    erlang:nif_error(undefined).
+
 %% @doc Associate a native code stream with a module
 %% @return ok
 %% @param Module module to set the native code of
@@ -154,10 +163,16 @@ load(Module) ->
                         LiteralResolver = fun(Index) ->
                             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),
                         {LabelsCount, BackendState1} = jit:compile(
-                            Code, AtomResolver, LiteralResolver, BackendModule, BackendState0
+                            Code,
+                            AtomResolver,
+                            LiteralResolver,
+                            TypeResolver,
+                            BackendModule,
+                            BackendState0
                         ),
                         Stream1 = BackendModule:stream(BackendState1),
                         code_server:set_native_code(Module, LabelsCount, Stream1),

libs/jit/src/jit.erl

@@ -24,7 +24,7 @@
     stream/1,
     backend/1,
     beam_chunk_header/3,
-    compile/5
+    compile/6
 ]).
 
 % NIFs
@@ -98,7 +98,8 @@
     line_offsets :: [{integer(), integer()}],
     labels_count :: pos_integer(),
     atom_resolver :: fun((integer()) -> atom()),
-    literal_resolver :: fun((integer()) -> any())
+    literal_resolver :: fun((integer()) -> any()),
+    type_resolver :: fun((integer()) -> any())
 }).
 
 -type stream() :: any().
@@ -130,6 +131,7 @@ compile(
     <<16:32, 0:32, OpcodeMax:32, LabelsCount:32, _FunctionsCount:32, Opcodes/binary>>,
     AtomResolver,
     LiteralResolver,
+    TypeResolver,
     MMod,
     MSt0
 ) when OpcodeMax =< ?OPCODE_MAX ->
@@ -138,7 +140,8 @@ compile(
         line_offsets = [],
         labels_count = LabelsCount,
         atom_resolver = AtomResolver,
-        literal_resolver = LiteralResolver
+        literal_resolver = LiteralResolver,
+        type_resolver = TypeResolver
     },
     {State1, MSt2} = first_pass(Opcodes, MMod, MSt1, State0),
     MSt3 = second_pass(MMod, MSt2, State1),
@@ -147,11 +150,12 @@ compile(
     <<16:32, 0:32, OpcodeMax:32, _LabelsCount:32, _FunctionsCount:32, _Opcodes/binary>>,
     _AtomResolver,
     _LiteralResolver,
+    _TypeResolver,
     _MMod,
     _MSt
 ) ->
     error(badarg, [OpcodeMax]);
-compile(CodeChunk, _AtomResolver, _LiteralResolver, _MMod, _MSt) ->
+compile(CodeChunk, _AtomResolver, _LiteralResolver, _TypeResolver, _MMod, _MSt) ->
     error(badarg, [CodeChunk]).
 
 % 1
@@ -1143,7 +1147,7 @@ first_pass(<<?OP_IS_FUNCTION2, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {Label, Rest1} = decode_label(Rest0),
     {MSt1, Arg1, Rest2} = decode_compact_term(Rest1, MMod, MSt0, State0),
-    {MSt2, ArityTerm, Rest3} = decode_compact_term(Rest2, MMod, MSt1, State0),
+    {MSt2, ArityTerm, Rest3} = decode_typed_compact_term(Rest2, MMod, MSt1, State0),
     ?TRACE("OP_IS_FUNCTION2 ~p,~p,~p\n", [Label, Arg1, ArityTerm]),
     {MSt3, FuncPtr} = term_is_boxed_with_tag_and_get_ptr(Label, Arg1, ?TERM_BOXED_FUN, MMod, MSt2),
     {MSt4, Arity} = term_to_int(ArityTerm, Label, MMod, MSt3),
@@ -1174,7 +1178,7 @@ first_pass(<<?OP_BS_GET_INTEGER2, Rest0/binary>>, MMod, MSt0, State0) ->
     {Fail, Rest1} = decode_label(Rest0),
     {MSt1, Src, Rest2} = decode_compact_term(Rest1, MMod, MSt0, State0),
     {_Live, Rest3} = decode_literal(Rest2),
-    {MSt2, Size, Rest4} = decode_compact_term(Rest3, MMod, MSt1, State0),
+    {MSt2, Size, Rest4} = decode_typed_compact_term(Rest3, MMod, MSt1, State0),
     {Unit, Rest5} = decode_literal(Rest4),
     {FlagsValue, Rest6} = decode_literal(Rest5),
     {MSt3, SrcReg} = MMod:move_to_native_register(MSt2, Src),
@@ -1213,7 +1217,7 @@ first_pass(<<?OP_BS_GET_FLOAT2, Rest0/binary>>, MMod, MSt0, State0) ->
     {Fail, Rest1} = decode_label(Rest0),
     {MSt1, Src, Rest2} = decode_compact_term(Rest1, MMod, MSt0, State0),
     {_Live, Rest3} = decode_literal(Rest2),
-    {MSt2, Size, Rest4} = decode_compact_term(Rest3, MMod, MSt1, State0),
+    {MSt2, Size, Rest4} = decode_typed_compact_term(Rest3, MMod, MSt1, State0),
     {Unit, Rest5} = decode_literal(Rest4),
     {FlagsValue, Rest6} = decode_literal(Rest5),
     {MSt3, SrcReg} = MMod:move_to_native_register(MSt2, Src),
@@ -1338,7 +1342,7 @@ first_pass(<<?OP_BS_SKIP_BITS2, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {Fail, Rest1} = decode_label(Rest0),
     {MSt1, Src, Rest2} = decode_compact_term(Rest1, MMod, MSt0, State0),
-    {MSt2, Size, Rest3} = decode_compact_term(Rest2, MMod, MSt1, State0),
+    {MSt2, Size, Rest3} = decode_typed_compact_term(Rest2, MMod, MSt1, State0),
     {Unit, Rest4} = decode_literal(Rest3),
     {_FlagsValue, Rest5} = decode_literal(Rest4),
     ?TRACE("OP_BS_SKIP_BITS2 ~p, ~p, ~p, ~p, ~p\n", [Fail, Src, Size, Unit, _FlagsValue]),
@@ -2071,7 +2075,7 @@ first_pass(<<?OP_BS_GET_POSITION, Rest0/binary>>, MMod, MSt0, State0) ->
 first_pass(<<?OP_BS_SET_POSITION, Rest0/binary>>, MMod, MSt0, State0) ->
     ?ASSERT_ALL_NATIVE_FREE(MSt0),
     {MSt1, Src, Rest1} = decode_compact_term(Rest0, MMod, MSt0, State0),
-    {MSt2, Pos, Rest2} = decode_compact_term(Rest1, MMod, MSt1, State0),
+    {MSt2, Pos, Rest2} = decode_typed_compact_term(Rest1, MMod, MSt1, State0),
     ?TRACE("OP_BS_SET_POSITION ~p, ~p\n", [Src, Pos]),
     {MSt3, MatchStateReg} = MMod:move_to_native_register(MSt2, Src),
     {MSt4, MatchStateRegPtr} = verify_is_match_state_and_get_ptr(MMod, MSt3, {free, MatchStateReg}),
@@ -2814,7 +2818,7 @@ first_pass_bs_match_integer(
     {_Live, Rest1} = decode_literal(Rest0),
     {Flags, Rest2} = decode_compile_time_literal(Rest1, State0),
     {MSt1, FlagsValue} = decode_flags_list(Flags, MMod, MSt0),
-    {MSt2, Size, Rest3} = decode_compact_term(Rest2, MMod, MSt0, State0),
+    {MSt2, Size, Rest3} = decode_typed_compact_term(Rest2, MMod, MSt0, State0),
     {Unit, Rest4} = decode_literal(Rest3),
     ?TRACE("{integer,~p,~p,~p, ", [Flags, Size, Unit]),
     {MSt3, SizeReg} = term_to_int(Size, 0, MMod, MSt1),
@@ -3185,6 +3189,14 @@ term_to_int(Term, _FailLabel, _MMod, MSt0) when is_integer(Term) ->
     {MSt0, Term bsr 4};
 term_to_int({literal, Val}, _FailLabel, _MMod, MSt0) when is_integer(Val) ->
     {MSt0, Val};
+% Optimized case: when we have type information showing this is an integer, skip the type check
+term_to_int({typed, Term, {t_integer, _Range}}, _FailLabel, MMod, MSt0) ->
+    {MSt1, Reg} = MMod:move_to_native_register(MSt0, Term),
+    MSt2 = MMod:shift_right(MSt1, Reg, 4),
+    {MSt2, Reg};
+term_to_int({typed, Term, _NonIntegerType}, FailLabel, MMod, MSt0) ->
+    % Type information shows it's not an integer, fall back to generic path
+    term_to_int(Term, FailLabel, MMod, MSt0);
 term_to_int(Term, FailLabel, MMod, MSt0) ->
     {MSt1, Reg} = MMod:move_to_native_register(MSt0, Term),
     MSt2 = cond_raise_badarg_or_jump_to_fail_label(
@@ -3357,6 +3369,17 @@ decode_compact_term(<<_Value:5, ?COMPACT_LITERAL:3, _Rest/binary>> = Binary, _MM
 decode_compact_term(Other, MMod, MSt, _State) ->
     decode_dest(Other, MMod, MSt).
 
+% Decode compact term with type information awareness
+decode_typed_compact_term(<<?COMPACT_EXTENDED_TYPED_REGISTER, Rest0/binary>>, MMod, MSt0, #state{
+    type_resolver = TypeResover
+}) ->
+    {MSt1, Dest, Rest1} = decode_dest(Rest0, MMod, MSt0),
+    {TypeIx, Rest2} = decode_literal(Rest1),
+    Type = TypeResover(TypeIx),
+    {MSt1, {typed, Dest, Type}, Rest2};
+decode_typed_compact_term(Other, MMod, MSt, State) ->
+    decode_compact_term(Other, MMod, MSt, State).
+
 skip_compact_term(<<_:4, ?COMPACT_INTEGER:4, _Rest/binary>> = Bin) ->
     {_Value, Rest} = decode_value64(Bin),
     Rest;

libs/jit/src/jit_precompile.erl

@@ -19,7 +19,7 @@
 %
 -module(jit_precompile).
 
--export([start/0, compile/3]).
+-export([start/0, compile/3, atom_resolver/1, type_resolver/1]).
 
 -include_lib("jit.hrl").
 
@@ -36,8 +36,7 @@ compile(Target, Dir, Path) ->
         FilteredChunks = lists:keydelete("Code", 1, FilteredChunks0),
         {"Code", CodeChunk} = lists:keyfind("Code", 1, InitialChunks),
         {"AtU8", AtomChunk} = lists:keyfind("AtU8", 1, InitialChunks),
-        Atoms = parse_atom_chunk(AtomChunk),
-        AtomResolver = fun(Index) -> lists:nth(Index, Atoms) end,
+        AtomResolver = atom_resolver(AtomChunk),
         LiteralsChunk =
             case lists:keyfind("LitU", 1, InitialChunks) of
                 {"LitU", LiteralsChunk0} ->
@@ -52,8 +51,16 @@ compile(Target, Dir, Path) ->
                             <<>>
                     end
             end,
-        Literals = parse_literals_chunk(LiteralsChunk),
-        LiteralResolver = fun(Index) -> lists:nth(Index + 1, Literals) end,
+        LiteralResolver = literal_resolver(LiteralsChunk),
+
+        TypesChunk =
+            case lists:keyfind("Type", 1, InitialChunks) of
+                {"Type", TypesChunk0} ->
+                    TypesChunk0;
+                false ->
+                    <<>>
+            end,
+        TypeResolver = type_resolver(TypesChunk),
 
         Stream0 = jit_stream_binary:new(0),
         <<16:32, 0:32, _OpcodeMax:32, LabelsCount:32, _FunctionsCount:32, _Opcodes/binary>> =
@@ -64,7 +71,7 @@ compile(Target, Dir, Path) ->
         Backend = list_to_atom("jit_" ++ Target),
         Stream2 = Backend:new(?JIT_VARIANT_PIC, jit_stream_binary, Stream1),
         {LabelsCount, Stream3} = jit:compile(
-            CodeChunk, AtomResolver, LiteralResolver, Backend, Stream2
+            CodeChunk, AtomResolver, LiteralResolver, TypeResolver, Backend, Stream2
         ),
         NativeCode = Backend:stream(Stream3),
         UpdatedChunks = FilteredChunks ++ [{"avmN", NativeCode}],
@@ -78,6 +85,10 @@ compile(Target, Dir, Path) ->
             io:format("Unimplemented opcode ~p (~s)\n", [Opcode, Path])
     end.
 
+atom_resolver(AtomChunk) ->
+    Atoms = parse_atom_chunk(AtomChunk),
+    fun(Index) -> lists:nth(Index, Atoms) end.
+
 parse_atom_chunk(<<AtomCount:32/signed, Rest/binary>>) ->
     if
         AtomCount < 0 ->
@@ -100,6 +111,10 @@ parse_atom_chunk_old_format(<<Size, Atom:Size/binary, Tail/binary>>, Acc) ->
 parse_atom_chunk_old_format(<<>>, Acc) ->
     lists:reverse(Acc).
 
+literal_resolver(LiteralsChunk) ->
+    Literals = parse_literals_chunk(LiteralsChunk),
+    fun(Index) -> lists:nth(Index + 1, Literals) end.
+
 parse_literals_chunk(<<TermsCount:32, Rest/binary>>) ->
     parse_literals_chunk0(TermsCount, Rest, []);
 parse_literals_chunk(<<>>) ->
@@ -110,3 +125,82 @@ parse_literals_chunk0(0, <<>>, Acc) ->
 parse_literals_chunk0(N, <<TermSize:32, TermBin:TermSize/binary, Rest/binary>>, Acc) ->
     Term = binary_to_term(TermBin),
     parse_literals_chunk0(N - 1, Rest, [Term | Acc]).
+
+%% Version (from beam_types.hrl)
+-define(BEAM_TYPES_VERSION, 3).
+
+%% Type chunk constants (from beam_types.erl)
+-define(BEAM_TYPE_ATOM, (1 bsl 0)).
+-define(BEAM_TYPE_BITSTRING, (1 bsl 1)).
+-define(BEAM_TYPE_CONS, (1 bsl 2)).
+-define(BEAM_TYPE_FLOAT, (1 bsl 3)).
+-define(BEAM_TYPE_FUN, (1 bsl 4)).
+-define(BEAM_TYPE_INTEGER, (1 bsl 5)).
+-define(BEAM_TYPE_MAP, (1 bsl 6)).
+-define(BEAM_TYPE_NIL, (1 bsl 7)).
+-define(BEAM_TYPE_PID, (1 bsl 8)).
+-define(BEAM_TYPE_PORT, (1 bsl 9)).
+-define(BEAM_TYPE_REFERENCE, (1 bsl 10)).
+-define(BEAM_TYPE_TUPLE, (1 bsl 11)).
+
+-define(BEAM_TYPE_HAS_LOWER_BOUND, (1 bsl 12)).
+-define(BEAM_TYPE_HAS_UPPER_BOUND, (1 bsl 13)).
+-define(BEAM_TYPE_HAS_UNIT, (1 bsl 14)).
+
+type_resolver(<<Version:32, _Count:32, TypeData/binary>>) when Version =:= ?BEAM_TYPES_VERSION ->
+    Types = parse_type_entries(TypeData, []),
+    fun(Index) -> lists:nth(Index + 1, Types) end;
+type_resolver(_) ->
+    fun(_) -> any end.
+
+parse_type_entries(<<>>, Acc) ->
+    lists:reverse(Acc);
+parse_type_entries(
+    <<0:1, HasUnit:1, HasUpperBound:1, HasLowerBound:1, TypeBits:12, Rest0/binary>>, Acc
+) ->
+    {Rest, LowerBound, UpperBound, Unit} = parse_extra(
+        HasLowerBound, HasUpperBound, HasUnit, Rest0, '-inf', '+inf', 1
+    ),
+    Type =
+        case TypeBits of
+            ?BEAM_TYPE_ATOM ->
+                t_atom;
+            ?BEAM_TYPE_BITSTRING ->
+                {t_bs_matchable, Unit};
+            ?BEAM_TYPE_CONS ->
+                t_cons;
+            ?BEAM_TYPE_FLOAT ->
+                t_float;
+            ?BEAM_TYPE_FUN ->
+                t_fun;
+            ?BEAM_TYPE_FLOAT bor ?BEAM_TYPE_INTEGER ->
+                {t_number, {LowerBound, UpperBound}};
+            ?BEAM_TYPE_INTEGER ->
+                {t_integer, {LowerBound, UpperBound}};
+            ?BEAM_TYPE_MAP ->
+                t_map;
+            ?BEAM_TYPE_NIL ->
+                nil;
+            ?BEAM_TYPE_NIL bor ?BEAM_TYPE_CONS ->
+                t_list;
+            ?BEAM_TYPE_PID ->
+                pid;
+            ?BEAM_TYPE_PORT ->
+                port;
+            ?BEAM_TYPE_REFERENCE ->
+                reference;
+            ?BEAM_TYPE_TUPLE ->
+                t_tuple;
+            _ ->
+                any
+        end,
+    parse_type_entries(Rest, [Type | Acc]).
+
+parse_extra(1, HasUpperBound, HasUnit, <<Value:64/signed, Rest/binary>>, '-inf', '+inf', 1) ->
+    parse_extra(0, HasUpperBound, HasUnit, Rest, Value, '+inf', 1);
+parse_extra(0, 1, HasUnit, <<Value:64/signed, Rest/binary>>, LowerBound, '+inf', 1) ->
+    parse_extra(0, 0, HasUnit, Rest, LowerBound, Value, 1);
+parse_extra(0, 0, 1, <<Value:8/unsigned, Rest/binary>>, LowerBound, UpperBound, 1) ->
+    parse_extra(0, 0, 0, Rest, LowerBound, UpperBound, Value + 1);
+parse_extra(0, 0, 0, Rest, LowerBound, UpperBound, Unit) ->
+    {Rest, LowerBound, UpperBound, Unit}.

src/libAtomVM/iff.c

@@ -100,6 +100,9 @@ void scan_iff(const void *iff_binary, int buf_size, unsigned long *offsets, unsi
         } else if (!memcmp(current_record->name, "avmN", 4)) {
             offsets[AVMN] = current_pos;
             sizes[AVMN] = ENDIAN_SWAP_32(current_record->size);
+        } else if (!memcmp(current_record->name, "Type", 4)) {
+            offsets[TYPE] = current_pos;
+            sizes[TYPE] = ENDIAN_SWAP_32(current_record->size);
         }
 
         current_pos += iff_align(ENDIAN_SWAP_32(current_record->size) + 8);

src/libAtomVM/iff.h

@@ -58,11 +58,13 @@ extern "C" {
 #define LINT 9
 /** Native code section */
 #define AVMN 10
+/** Type table section */
+#define TYPE 11
 
 /** Required size for offsets array */
-#define MAX_OFFS 11
+#define MAX_OFFS 12
 /** Required size for sizes array */
-#define MAX_SIZES 11
+#define MAX_SIZES 12
 
 /** sizeof IFF section header in bytes */
 #define IFF_SECTION_HEADER_SIZE 8