JIT: make imported functions available to compiler

Optimize + and - as well as < and > when types are integers of known range.

Also fix array functions of armv6m backend for higher indexes (>= 32).

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

Author: pguyot

libs/estdlib/src/code_server.erl

@@ -39,6 +39,7 @@
     atom_resolver/2,
     literal_resolver/2,
     type_resolver/2,
+    import_resolver/2,
     set_native_code/3
 ]).
 
@@ -135,6 +136,14 @@ literal_resolver(_Module, _Index) ->
 type_resolver(_Module, _Index) ->
     erlang:nif_error(undefined).
 
+%% @doc Get an imported function triplet from its index
+%% @return The imported function as {Module, Function, Arity}
+%% @param Module module to get the imported function from
+%% @param Index imported function index in the module
+-spec import_resolver(Module :: module(), Index :: non_neg_integer()) -> {atom(), atom(), non_neg_integer()}.
+import_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
@@ -164,13 +173,17 @@ load(Module) ->
                             code_server:literal_resolver(Module, Index)
                         end,
                         TypeResolver = fun(Index) -> code_server:type_resolver(Module, Index) end,
+                        ImportResolver = fun(Index) ->
+                            code_server:import_resolver(Module, Index)
+                        end,
                         {StreamModule, Stream0} = jit:stream(jit_mmap_size(byte_size(Code))),
                         {BackendModule, BackendState0} = jit:backend(StreamModule, Stream0),
                         {LabelsCount, BackendState1} = jit:compile(
                             Code,
                             AtomResolver,
                             LiteralResolver,
                             TypeResolver,
+                            ImportResolver,
                             BackendModule,
                             BackendState0
                         ),

libs/jit/src/jit.erl

@@ -154,6 +154,7 @@
 -type condition() ::
     {aarch64_register(), '<', integer()}
     | {maybe_free_aarch64_register(), '<', aarch64_register()}
+    | {integer(), '<', maybe_free_aarch64_register()}
     | {maybe_free_aarch64_register(), '==', integer()}
     | {maybe_free_aarch64_register(), '!=', aarch64_register() | integer()}
     | {'(int)', maybe_free_aarch64_register(), '==', integer()}
@@ -757,15 +758,47 @@ if_else_block(
         jit_aarch64_asm:cc() | {tbz | tbnz, atom(), 0..63} | {cbz, atom()},
         non_neg_integer()
     }.
-if_block_cond(#state{stream_module = StreamModule, stream = Stream0} = State0, {Reg, '<', 0}) ->
+if_block_cond(#state{stream_module = StreamModule, stream = Stream0} = State0, {RegOrTuple, '<', 0}) ->
+    Reg =
+        case RegOrTuple of
+            {free, Reg0} -> Reg0;
+            RegOrTuple -> RegOrTuple
+        end,
     I = jit_aarch64_asm:tbz(Reg, 63, 0),
     Stream1 = StreamModule:append(Stream0, I),
-    State1 = State0#state{stream = Stream1},
-    {State1, {tbz, Reg, 63}, 0};
+    State1 = if_block_free_reg(RegOrTuple, State0),
+    State2 = State1#state{stream = Stream1},
+    {State2, {tbz, Reg, 63}, 0};
+% Handle {Val, '<', Reg} - means Val < Reg, jump if false (i.e., if Val >= Reg or Reg <= Val)
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0,
-    {Reg, '<', Val}
-) when is_atom(Reg), is_integer(Val) ->
+    {Val, '<', RegOrTuple}
+) when is_integer(Val) ->
+    Reg =
+        case RegOrTuple of
+            {free, Reg0} -> Reg0;
+            RegOrTuple -> RegOrTuple
+        end,
+    I1 = jit_aarch64_asm:cmp(Reg, Val),
+    % le = less than or equal
+    I2 = jit_aarch64_asm:bcc(le, 0),
+    Code = <<
+        I1/binary,
+        I2/binary
+    >>,
+    Stream1 = StreamModule:append(Stream0, Code),
+    State1 = if_block_free_reg(RegOrTuple, State0),
+    State2 = State1#state{stream = Stream1},
+    {State2, le, byte_size(I1)};
+if_block_cond(
+    #state{stream_module = StreamModule, stream = Stream0} = State0,
+    {RegOrTuple, '<', Val}
+) when is_integer(Val), Val =/= 0 ->
+    Reg =
+        case RegOrTuple of
+            {free, Reg0} -> Reg0;
+            RegOrTuple -> RegOrTuple
+        end,
     I1 = jit_aarch64_asm:cmp(Reg, Val),
     % ge = greater than or equal
     I2 = jit_aarch64_asm:bcc(ge, 0),
@@ -774,8 +807,9 @@ if_block_cond(
         I2/binary
     >>,
     Stream1 = StreamModule:append(Stream0, Code),
-    State1 = State0#state{stream = Stream1},
-    {State1, ge, byte_size(I1)};
+    State1 = if_block_free_reg(RegOrTuple, State0),
+    State2 = State1#state{stream = Stream1},
+    {State2, ge, byte_size(I1)};
 if_block_cond(
     #state{stream_module = StreamModule, stream = Stream0} = State0,
     {RegOrTuple, '<', RegB}
@@ -2073,8 +2107,10 @@ add(State, Reg, Val) ->
 %% @param Val immediate value to subtract
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
--spec sub(state(), aarch64_register(), integer()) -> state().
-sub(#state{stream_module = StreamModule, stream = Stream0} = State, Reg, Val) ->
+-spec sub(state(), aarch64_register(), integer() | aarch64_register()) -> state().
+sub(State, Reg, Val) when is_integer(Val) ->
+    op_imm(State, sub, Reg, Reg, Val);
+sub(#state{stream_module = StreamModule, stream = Stream0} = State, Reg, Val) when is_atom(Val) ->
     I1 = jit_aarch64_asm:sub(Reg, Reg, Val),
     Stream1 = StreamModule:append(Stream0, I1),
     State#state{stream = Stream1}.

libs/jit/src/jit_aarch64.erl

@@ -948,6 +948,8 @@ sub(Rd, Rn, Imm) when is_atom(Rd), is_atom(Rn), is_integer(Imm), Imm >= 0, Imm =
     RdNum = reg_to_num(Rd),
     RnNum = reg_to_num(Rn),
     <<(16#D1000000 bor ((Imm band 16#FFF) bsl 10) bor (RnNum bsl 5) bor RdNum):32/little>>;
+sub(_Rd, _Rn, Imm) when is_integer(Imm) ->
+    error({unencodable_immediate, Imm});
 sub(Rd, Rn, Rm) when is_atom(Rd), is_atom(Rn), is_atom(Rm) ->
     sub(Rd, Rn, Rm, {lsl, 0}).